#include "vm_types.h"
#include <stdlib.h>
#include <string.h>
#include <stdio.h>
#include <assert.h>

// Safe memory allocation with overflow checking
static void* safe_malloc(size_t size) {
    if (size == 0 || size > SIZE_MAX / 2) return NULL;
    void* ptr = malloc(size);
    if (!ptr) {
        fprintf(stderr, "Memory allocation failed\n");
        exit(EXIT_FAILURE);
    }
    return ptr;
}

static void* safe_realloc(void* ptr, size_t size) {
    if (size == 0 || size > SIZE_MAX / 2) return NULL;
    void* new_ptr = realloc(ptr, size);
    if (!new_ptr) {
        fprintf(stderr, "Memory reallocation failed\n");
        exit(EXIT_FAILURE);
    }
    return new_ptr;
}

// Initialize VM with security limits
vm_t* vm_create(uint8_t* bytecode, size_t size) {
    vm_t* vm = safe_malloc(sizeof(vm_t));
    memset(vm, 0, sizeof(vm_t));

    vm->bytecode = bytecode;
    vm->bytecode_size = size;
    vm->ip = 0;
    vm->halted = false;

    // Security limits
    vm->max_stack_size = 65536;
    vm->max_call_depth = 256;
    vm->current_call_depth = 0;

    return vm;
}

// Safe bytecode reading with bounds checking
static bool read_u8(vm_t* vm, uint8_t* result) {
    if (vm->ip >= vm->bytecode_size) return false;
    *result = vm->bytecode[vm->ip++];
    return true;
}

static bool read_u16(vm_t* vm, uint16_t* result) {
    if (vm->ip + 2 > vm->bytecode_size) return false;
    memcpy(result, &vm->bytecode[vm->ip], 2);
    vm->ip += 2;
    return true;
}

static bool read_u32(vm_t* vm, uint32_t* result) {
    if (vm->ip + 4 > vm->bytecode_size) return false;
    memcpy(result, &vm->bytecode[vm->ip], 4);
    vm->ip += 4;
    return true;
}

static bool read_i32(vm_t* vm, int32_t* result) {
    if (vm->ip + 4 > vm->bytecode_size) return false;
    memcpy(result, &vm->bytecode[vm->ip], 4);
    vm->ip += 4;
    return true;
}

static bool read_f32(vm_t* vm, float* result) {
    if (vm->ip + 4 > vm->bytecode_size) return false;
    memcpy(result, &vm->bytecode[vm->ip], 4);
    vm->ip += 4;
    return true;
}

// Stack operations with bounds checking
static bool stack_push(frame_t* frame, value_t value) {
    if (frame->stack_size >= frame->stack_capacity) {
        size_t new_capacity = frame->stack_capacity * 2;
        if (new_capacity == 0) new_capacity = 16;

        value_t* new_stack = safe_realloc(frame->stack, new_capacity * sizeof(value_t));
        if (!new_stack) return false;

        frame->stack = new_stack;
        frame->stack_capacity = new_capacity;
    }

    frame->stack[frame->stack_size++] = value;
    return true;
}

static bool stack_pop(frame_t* frame, value_t* result) {
    if (frame->stack_size == 0) return false;
    *result = frame->stack[--frame->stack_size];
    return true;
}

// Type checking and conversion
static bool type_check_binary(value_t a, value_t b, type_code_t* result_type) {
    // Simple type checking - in real implementation, add proper type promotion
    if (a.type == b.type) {
        *result_type = a.type;
        return true;
    }
    return false;
}

static int32_t value_to_int(value_t val) {
    switch (val.type) {
    case TYPE_I8: return val.data.i8;
    case TYPE_U8: return val.data.u8;
    case TYPE_I16: return val.data.i16;
    case TYPE_U16: return val.data.u16;
    case TYPE_I32: return val.data.i32;
    case TYPE_U32: return (int32_t)val.data.u32;
    case TYPE_BOOL: return val.data.boolean ? 1 : 0;
    case TYPE_CHAR: return (int32_t)val.data.character;
    default: return 0;
    }
}

static float value_to_float(value_t val) {
    switch (val.type) {
    case TYPE_F32: return val.data.f32;
    case TYPE_I32: return (float)val.data.i32;
    case TYPE_U32: return (float)val.data.u32;
    default: return 0.0f;
    }
}

// Bytecode validation
bool vm_validate_bytecode(vm_t* vm) {
    // Check magic number
    if (vm->bytecode_size < 4 || memcmp(vm->bytecode, "POPC", 4) != 0) {
        return false;
    }

    // Add more validation as needed
    return true;
}

// Execute single instruction with full error checking
bool vm_execute_instruction(vm_t* vm) {
    if (vm->halted || vm->ip >= vm->bytecode_size) return false;

    uint8_t opcode;
    if (!read_u8(vm, &opcode)) return false;

    frame_t* frame = vm->current_frame;
    if (!frame) return false;

    switch (opcode) {
    case OP_PUSH_CONST: {
        uint32_t const_idx;
        if (!read_u32(vm, &const_idx) || const_idx >= vm->constants_count) {
            return false;
        }
        if (!stack_push(frame, vm->constants[const_idx])) return false;
        break;
    }

    case OP_PUSH_INT: {
        uint8_t width;
        int32_t value;
        if (!read_u8(vm, &width) || !read_i32(vm, &value)) return false;

        value_t val = { 0 };
        switch (width) {
        case 8: val.type = TYPE_I8; val.data.i8 = (int8_t)value; break;
        case 16: val.type = TYPE_I16; val.data.i16 = (int16_t)value; break;
        case 32: val.type = TYPE_I32; val.data.i32 = value; break;
        default: return false;
        }
        if (!stack_push(frame, val)) return false;
        break;
    }

    case OP_ADD: {
        value_t a, b;
        if (!stack_pop(frame, &a) || !stack_pop(frame, &b)) return false;

        type_code_t result_type;
        if (!type_check_binary(a, b, &result_type)) return false;

        value_t result = { 0 };
        result.type = result_type;

        if (result_type == TYPE_F32) {
            result.data.f32 = value_to_float(b) + value_to_float(a);
        }
        else {
            result.data.i32 = value_to_int(b) + value_to_int(a);
        }

        if (!stack_push(frame, result)) return false;
        break;
    }

    case OP_CMP_EQ: {
        value_t a, b;
        if (!stack_pop(frame, &a) || !stack_pop(frame, &b)) return false;

        value_t result = { 0 };
        result.type = TYPE_BOOL;
        result.data.boolean = (value_to_int(a) == value_to_int(b));

        if (!stack_push(frame, result)) return false;
        break;
    }

    case OP_JMP: {
        int32_t offset;
        if (!read_i32(vm, &offset)) return false;

        // Check for negative jumps (security)
        if (offset < 0 && (size_t)(-offset) > vm->ip) return false;
        if (vm->ip + offset >= vm->bytecode_size) return false;

        vm->ip += offset;
        break;
    }

    case OP_PRINT: {
        value_t val;
        if (!stack_pop(frame, &val)) return false;

        switch (val.type) {
        case TYPE_I32: printf("%d\n", val.data.i32); break;
        case TYPE_F32: printf("%f\n", val.data.f32); break;
        case TYPE_BOOL: printf("%s\n", val.data.boolean ? "true" : "false"); break;
        case TYPE_STR: printf("%s\n", val.data.string); break;
        default: printf("<unknown>\n"); break;
        }
        break;
    }

    case OP_HALT:
        vm->halted = true;
        break;

    default:
        fprintf(stderr, "Unknown opcode: 0x%02x\n", opcode);
        return false;
    }

    return true;
}

// Main execution loop
void vm_execute(vm_t* vm) {
    while (!vm->halted && vm_execute_instruction(vm)) {
        // Continue execution
    }
}

// Cleanup
void vm_destroy(vm_t* vm) {
    if (!vm) return;

    // Free constants
    for (size_t i = 0; i < vm->constants_count; i++) {
        if (vm->constants[i].type == TYPE_STR && vm->constants[i].data.string) {
            free(vm->constants[i].data.string);
        }
    }
    free(vm->constants);

    // Free functions
    free(vm->functions);

    // Free frames (simplified - in real impl, walk call stack)
    if (vm->current_frame) {
        free(vm->current_frame->locals);
        free(vm->current_frame->stack);
        free(vm->current_frame);
    }

    free(vm);
}