#include "vm.h"
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <stdint.h>

const char* vm_error_to_string(VMError error) {
    switch (error) {
        case VM_SUCCESS: return "No error";
        case VM_ERROR_STACK_OVERFLOW: return "Stack overflow";
        case VM_ERROR_STACK_UNDERFLOW: return "Stack underflow";
        case VM_ERROR_INVALID_OPCODE: return "Invalid opcode";
        case VM_ERROR_INVALID_REGISTER: return "Invalid register access";
        case VM_ERROR_INVALID_CONSTANT: return "Invalid constant index";
        case VM_ERROR_INVALID_JUMP: return "Invalid jump address";
        case VM_ERROR_INVALID_NATIVE_FUNCTION: return "Invalid native function";
        case VM_ERROR_MEMORY_ALLOCATION: return "Memory allocation failed";
        case VM_ERROR_DIVISION_BY_ZERO: return "Division by zero";
        case VM_ERROR_INVALID_FRAME: return "Invalid frame operation";
        default: return "Unknown error";
    }
}

void vm_set_error(VM* vm, VMError error, const char* message, uint64_t ip, uint8_t opcode) {
    if (!vm) return;
    vm->error = error;
    vm->error_ip = ip;
    vm->error_opcode = opcode;
    
    // 格式化错误消息，包含位置和操作码信息
    char full_message[256];
    snprintf(full_message, sizeof(full_message), 
             "[IP: %lu, Opcode: 0x%02X] %s", 
             ip, opcode, message ? message : vm_error_to_string(error));
    
    strncpy(vm->error_message, full_message, sizeof(vm->error_message) - 1);
    vm->error_message[sizeof(vm->error_message) - 1] = '\0';
}

VMError vm_get_error(VM* vm) {
    return vm ? vm->error : VM_ERROR_INVALID_REGISTER;
}

const char* vm_get_error_message(VM* vm) {
    return vm ? vm->error_message : "Invalid VM instance";
}

VM* vm_create() {
    VM* vm = (VM*)malloc(sizeof(VM));
    if (!vm) return NULL;

    // 初始化寄存器
    memset(vm->registers, 0, sizeof(uint64_t) * REGISTER_COUNT);
    
    // 初始化栈
    vm->sp = 0;
    vm->ip = 0;
    vm->frame_count = 0;
    vm->frames = NULL;
    
    // 初始化程序
    vm->program = NULL;
    
    // 初始化全局变量表
    vm->global_count = 0;
    memset(vm->globals, 0, sizeof(GlobalVar) * GLOBAL_VAR_COUNT);

    // 初始化函数表
    vm->native_functions = malloc(sizeof(NativeFunction) * 10); // 增加初始容量
    vm->native_function_count = 0;
    vm->functions = NULL;
    vm->function_count = 0;

    // 初始化错误状态
    vm->error = VM_SUCCESS;
    vm->error_ip = 0;
    vm->error_opcode = 0;
    memset(vm->error_message, 0, sizeof(vm->error_message));
    
    vm->block_count = 0;

    return vm;
}

void vm_destroy(VM* vm) {
    if (!vm) return;
    
    // 释放程序内存
    if (vm->program) {
        if (vm->program->constants.entries) {
            // 释放常量池中的字符串
            for (uint64_t i = 0; i < vm->program->constants.count; i++) {
                if (vm->program->constants.entries[i].type == TYPE_STRING) {
                    //free(vm->program->constants.entries[i].value.string_val.value);
                }
            }
            //free(vm->program->constants.entries);
        }
        //free(vm->program);
    }
    
    // 释放全局变量表中的字符串
    for (size_t i = 0; i < vm->global_count; i++) {
        free(vm->globals[i].name);
        if (vm->globals[i].type == TYPE_STRING) {
            free(vm->globals[i].value.string_val.value);
        }
    }
    
    // 释放函数表
    if (vm->native_functions) {
        for (size_t i = 0; i < vm->native_function_count; i++) {
            free(vm->native_functions[i].name);
        }
        free(vm->native_functions);
    }

    if (vm->functions) {
        for (size_t i = 0; i < vm->function_count; i++) {
            free(vm->functions[i].name);
        }
        free(vm->functions);
    }
    
    // 释放栈帧
    free(vm->frames);
    free(vm);
}

void vm_set_register(VM* vm, uint64_t index, uint64_t value) {
    if (!vm || index >= REGISTER_COUNT) {
        // 无效的寄存器访问
        fprintf(stderr, "Invalid register access\n");
        return;
    }
    vm->registers[index] = value;
}

uint64_t vm_get_register(VM* vm, uint64_t index) {
    if (!vm || index >= REGISTER_COUNT) {
        // 无效的寄存器访问
        fprintf(stderr, "Invalid register access\n");
        return 0;
    }
    return vm->registers[index];
}

void vm_push(VM* vm, uint64_t value) {
    if (!vm || vm->sp >= STACK_SIZE) {
        // 栈溢出处理
        fprintf(stderr, "Stack overflow\n");
        return;
    }
    vm->stack[vm->sp++] = value;
}

uint64_t vm_pop(VM* vm) {
    if (!vm || vm->sp == 0) {
        // 栈下溢处理
        fprintf(stderr, "Stack underflow\n");
        return 0;
    }
    return vm->stack[--vm->sp];
}

uint64_t vm_peek(VM* vm, uint64_t offset) {
    if (!vm || vm->sp <= offset) {
        // 无效的栈访问
        fprintf(stderr, "Invalid stack access\n");
        return 0;
    }
    return vm->stack[vm->sp - 1 - offset];
}

int vm_register_native_function(VM* vm, const char* name, Value (*func)(VM* vm, Value* args, uint8_t argc)) {
    if (!vm || !name || !func) {
        fprintf(stderr, "Invalid parameters for native function registration\n");
        return -1;
    }
    
    // 检查函数是否已经注册
    for (size_t i = 0; i < vm->native_function_count; i++) {
        if (strcmp(vm->native_functions[i].name, name) == 0) {
            // 函数已存在，更新函数指针
            vm->native_functions[i].func = func;
            return i;
        }
    }
    
    // 如果需要，扩展原生函数表
    if (vm->native_function_count % 10 == 0 && vm->native_function_count > 0) {
        size_t new_size = (vm->native_function_count + 10) * sizeof(NativeFunction);
        NativeFunction* new_functions = realloc(vm->native_functions, new_size);
        if (!new_functions) {
            fprintf(stderr, "Failed to allocate memory for native functions\n");
            return -1;
        }
        vm->native_functions = new_functions;
    }
    
    // 添加新函数
    size_t idx = vm->native_function_count;
    vm->native_functions[idx].name = strdup(name);
    vm->native_functions[idx].func = func;
    vm->native_function_count++;
    
    return idx;
}

void vm_load(VM* vm, Program *program) {
    if (!vm || !program) return;
    vm->program = program;
    
    // 重置虚拟机状态
    vm->ip = 0;
    vm->sp = 0;
    vm->frame_count = 0;
    
    // 初始化全局变量表
    vm->global_count = 0;
    memset(vm->globals, 0, sizeof(GlobalVar) * GLOBAL_VAR_COUNT);

    // 初始化函数表
    vm->functions = NULL;
    vm->function_count = 0;
}


void vm_eval(VM* vm) {
    vm->error = VM_SUCCESS;
    if (!vm || !vm->program) {
        vm_set_error(vm, VM_ERROR_INVALID_REGISTER, "Invalid VM or no program loaded", 0, 0);
        return;
    }
    
    uint8_t dst_reg = 0;
    uint8_t src_reg1 = 0;
    uint8_t src_reg2 = 0;
    uint8_t cond_reg = 0;
    
    while (vm->ip < vm->program->code_size) {
        uint8_t opcode = vm->program->code[vm->ip++];
        
        switch (opcode) {
            case OP_LOAD_CONST: {
                uint8_t dst_reg = vm->program->code[vm->ip++];
                uint8_t const_idx = vm->program->code[vm->ip++];
                if (const_idx >= vm->program->constants.count) {
                    vm_set_error(vm, VM_ERROR_INVALID_CONSTANT, 
                               "Invalid constant index", 
                               vm->ip - 3, opcode);
                    return;
                }
                Constant constant = vm->program->constants.entries[const_idx];
                switch (constant.type) {
                    case TYPE_INT8:
                        vm->registers[dst_reg] = constant.value.int8_val;
                        break;
                    case TYPE_INT16:
                        vm->registers[dst_reg] = constant.value.int16_val;
                        break;
                    case TYPE_INT32:
                        vm->registers[dst_reg] = constant.value.int32_val;
                        break;
                    case TYPE_INT64:
                        vm->registers[dst_reg] = constant.value.int64_val;
                        break;
                    case TYPE_BOOL:
                        vm->registers[dst_reg] = constant.value.bool_val;
                        break;
                    case TYPE_CHAR:
                        vm->registers[dst_reg] = constant.value.char_val;
                        break;
                    default:
                        vm_set_error(vm, VM_ERROR_INVALID_CONSTANT,
                                   "Unsupported constant type",
                                   vm->ip - 3, opcode);
                        return;
                }
                break;
            }
            case OP_GLOBAL_LOAD: {
                uint8_t dst_reg = vm->program->code[vm->ip++];
                uint8_t global_idx = vm->program->code[vm->ip++];
                if (global_idx >= vm->global_count) {
                    fprintf(stderr, "Invalid global variable index: %d\n", global_idx);
                    return;
                }
                vm->registers[dst_reg] = vm->globals[global_idx].value.int32_val;
                break;
            }
            case OP_GLOBAL_STORE: {
                uint8_t src_reg = vm->program->code[vm->ip++];
                uint8_t global_idx = vm->program->code[vm->ip++];
                if (global_idx >= vm->global_count) {
                    fprintf(stderr, "Invalid global variable index: %d\n", global_idx);
                    return;
                }
                vm->globals[global_idx].value.int32_val = vm->registers[src_reg];
                break;
            }
            case OP_LOAD_LOCAL: {
                if (vm->frame_count == 0) {
                    fprintf(stderr, "No active frame for local variable access\n");
                    return;
                }
                uint8_t dst_reg = vm->program->code[vm->ip++];
                uint8_t local_idx = vm->program->code[vm->ip++];
                Frame* frame = &vm->frames[vm->frame_count - 1];
                vm->registers[dst_reg] = frame->locals[local_idx];
                break;
            }
            case OP_STORE_LOCAL: {
                if (vm->frame_count == 0) {
                    fprintf(stderr, "No active frame for local variable access\n");
                    return;
                }
                uint8_t src_reg = vm->program->code[vm->ip++];
                uint8_t local_idx = vm->program->code[vm->ip++];
                Frame* frame = &vm->frames[vm->frame_count - 1];
                frame->locals[local_idx] = vm->registers[src_reg];
                break;
            }
            case OP_FUNC_DEF: {
                uint8_t func_name_idx = vm->program->code[vm->ip++];
                uint8_t param_count = vm->program->code[vm->ip++];
                uint8_t local_count = vm->program->code[vm->ip++];
                uint16_t instruction_count = (vm->program->code[vm->ip] << 8) | vm->program->code[vm->ip + 1];
                vm->ip += 2; // Skip instruction count bytes

                // 扩展函数表
                size_t new_size = (vm->function_count + 1) * sizeof(Function);
                Function* new_functions = realloc(vm->functions, new_size);
                if (!new_functions) {
                    fprintf(stderr, "Failed to allocate memory for function\n");
                    return;
                }
                vm->functions = new_functions;

                Function* func = &vm->functions[vm->function_count++];
                func->name = strdup(vm->program->constants.entries[func_name_idx].value.string_val.value);
                func->code_offset = vm->ip;
                func->param_count = param_count;
                func->local_count = local_count;
                func->instructions_count = instruction_count;

                // 跳过函数体
                vm->ip += instruction_count;
                break;
            }
            case OP_CALL: {
                uint8_t func_idx = vm->program->code[vm->ip++];
                if (func_idx >= vm->function_count) {
                    fprintf(stderr, "Invalid function index: %d\n", func_idx);
                    return;
                }

                Function* func = &vm->functions[func_idx];

                // 创建新的栈帧
                vm->frames = realloc(vm->frames, (vm->frame_count + 1) * sizeof(Frame));
                if (!vm->frames) {
                    fprintf(stderr, "Failed to allocate memory for frame\n");
                    return;
                }

                Frame* frame = &vm->frames[vm->frame_count++];
                frame->ip = vm->ip;
                frame->locals = calloc(func->local_count, sizeof(uint64_t));
                frame->stack = calloc(STACK_FRAME_SIZE, sizeof(uint64_t));
                frame->sp = 0;

                // 保存当前IP并跳转到函数代码
                vm->ip = func->code_offset;
                break;
            }
        
            case OP_RETURN: {
                if (vm->frame_count == 0) {
                    fprintf(stderr, "No frame to return from\n");
                    return;
                }
                
                // 恢复前一个栈帧
                Frame* frame = &vm->frames[--vm->frame_count];
                vm->ip = frame->ip;
                
                // 释放栈帧资源
                free(frame->locals);
                free(frame->stack);
                break;
            }
            case OP_ADD: {
                dst_reg = vm->program->code[vm->ip++];
                uint8_t src_reg1 = vm->program->code[vm->ip++];
                uint8_t src_reg2 = vm->program->code[vm->ip++];
                vm->registers[dst_reg] = vm->registers[src_reg1] + vm->registers[src_reg2];
                break;
            }
            case OP_JMP: {
                uint16_t offset = (vm->program->code[vm->ip] << 8) | vm->program->code[vm->ip + 1];
                if (offset >= vm->program->code_size) {
                    vm_set_error(vm, VM_ERROR_INVALID_JUMP, 
                               "Invalid jump address",
                               vm->ip, opcode);
                    return;
                }
                vm->ip = offset;
                vm->ip += 2;  // 跳过偏移量字节
                break;
            }
            case OP_JMP_IF_FALSE: {
                cond_reg = vm->program->code[vm->ip++];
                uint16_t offset = (vm->program->code[vm->ip] << 8) | vm->program->code[vm->ip + 1];
                vm->ip += 2;
                
                if (!vm->registers[cond_reg]) {
                    // 条件为假时，跳转到指定位置
                    if (offset >= vm->program->code_size) {
                        vm_set_error(vm, VM_ERROR_INVALID_JUMP, 
                                   "Invalid jump address",
                                   vm->ip - 3, opcode);
                        return;
                    }
                    vm->ip = offset;
                    
                    // 跳过此块中的所有嵌套if，直到找到匹配的END_IF
                    vm->block_count++;
                }
                break;
            }
            
            case OP_END_IF: {
                // 处理if块的结束
                if (vm->block_count > 0) {
                    vm->block_count--;
                }
                break;
            }
            
            case OP_JMP_IF_TRUE: {
                cond_reg = vm->program->code[vm->ip++];
                uint16_t offset = (vm->program->code[vm->ip] << 8) | vm->program->code[vm->ip + 1];
                if (vm->registers[cond_reg]) {
                    vm->ip = offset;
                } else {
                    vm->ip += 2;
                }
                break;
            }
            case OP_CMP_EQ: {
                dst_reg = vm->program->code[vm->ip++];
                uint8_t src_reg1 = vm->program->code[vm->ip++];
                uint8_t src_reg2 = vm->program->code[vm->ip++];
                vm->registers[dst_reg] = vm->registers[src_reg1] == vm->registers[src_reg2];
                break;
            }
            case OP_CMP_NE: {
                dst_reg = vm->program->code[vm->ip++];
                uint8_t src_reg1 = vm->program->code[vm->ip++];
                uint8_t src_reg2 = vm->program->code[vm->ip++];
                vm->registers[dst_reg] = vm->registers[src_reg1] != vm->registers[src_reg2];
                break;
            }
            case OP_CMP_LT: {
                dst_reg = vm->program->code[vm->ip++];
                uint8_t src_reg1 = vm->program->code[vm->ip++];
                uint8_t src_reg2 = vm->program->code[vm->ip++];
                vm->registers[dst_reg] = vm->registers[src_reg1] < vm->registers[src_reg2];
                break;
            }
            case OP_CMP_LE: {
                dst_reg = vm->program->code[vm->ip++];
                uint8_t src_reg1 = vm->program->code[vm->ip++];
                uint8_t src_reg2 = vm->program->code[vm->ip++];
                vm->registers[dst_reg] = vm->registers[src_reg1] <= vm->registers[src_reg2];
                break;
            }
            case OP_CMP_GT: {
                dst_reg = vm->program->code[vm->ip++];
                uint8_t src_reg1 = vm->program->code[vm->ip++];
                uint8_t src_reg2 = vm->program->code[vm->ip++];
                vm->registers[dst_reg] = vm->registers[src_reg1] > vm->registers[src_reg2];
                break;
            }
            case OP_CMP_GE: {
                dst_reg = vm->program->code[vm->ip++];
                uint8_t src_reg1 = vm->program->code[vm->ip++];
                uint8_t src_reg2 = vm->program->code[vm->ip++];
                vm->registers[dst_reg] = vm->registers[src_reg1] >= vm->registers[src_reg2];
                break;
            }
            case OP_NATIVE_CALL: {
                uint8_t func_idx = vm->program->code[vm->ip++];
                uint8_t argc = vm->program->code[vm->ip++];
                
                if (func_idx >= vm->native_function_count) {
                    vm_set_error(vm, VM_ERROR_INVALID_NATIVE_FUNCTION, 
                               "Invalid native function index",
                               vm->ip - 2, opcode);
                    return;
                }
                
                // 准备参数数组
                Value args[argc];
                for (int i = 0; i < argc; i++) {
                    uint8_t arg_reg = vm->program->code[vm->ip++];
                    args[i].int64_val = vm->registers[arg_reg];
                }
                
                // 调用原生函数
                NativeFunction* func = &vm->native_functions[func_idx];
                Value result = func->func(vm, args, argc);
                vm->registers[R0] = result.int64_val;
                break;
            }
            case OP_HALT:
                return;
            case OP_MOV: {
                uint8_t dst_reg = vm->program->code[vm->ip++];
                uint8_t src_reg = vm->program->code[vm->ip++];
                vm->registers[dst_reg] = vm->registers[src_reg];
                break;
            }
            case OP_MOV_IMM: {
                uint8_t dst_reg = vm->program->code[vm->ip++];
                uint8_t imm = vm->program->code[vm->ip++];
                vm->registers[dst_reg] = imm;
                break;
            }
            default:
                vm_set_error(vm, VM_ERROR_INVALID_OPCODE, 
                           "Unknown instruction",
                           vm->ip - 1, opcode);
                return;
        }
        
        // 检查每条指令执行后是否有错误
        if (vm->error != VM_SUCCESS) {
            return;
        }
    }
}

void print_bytecode(Program program) {
    printf("字节码指令:\n");
    printf("地址   |   操作码   |   操作数\n");
    printf("-------------------------------------------\n");
    int indent = 0;
    for (size_t i = 0; i < program.code_size;) {
        for (int j = 0; j < indent; j++) printf("  ");
        printf("0x%04d   |   ", i);

        uint8_t opcode = program.code[i++];
        switch (opcode) {
            case OP_BLOCK_START:
                printf("BLOCK_START\n");
                indent++;
                break;
            case OP_BLOCK_END:
                indent--;
                printf("BLOCK_END\n");
                break;
            case OP_LOAD_CONST:
                printf("LOAD_CONST    R%d, %d\n", program.code[i], program.code[i+1]);
                i += 2;
                break;
            case OP_MOV:
                printf("MOV          R%d, %d\n", program.code[i], program.code[i+1]);
                i += 2;
                break;
            case OP_MOV_IMM:
                printf("MOV_IMM      R%d, %d\n", program.code[i], program.code[i+1]);
                i += 2;
                break;
            case OP_ADD:
                printf("ADD          R%d, R%d, R%d\n",
                       program.code[i], program.code[i+1], program.code[i+2]);
                i += 3;
                break;
            case OP_CMP_EQ:
                printf("CMP_EQ       R%d, R%d, R%d\n",
                       program.code[i], program.code[i+1], program.code[i+2]);
                i += 3;
                break;
            case OP_JMP_IF_FALSE:
                printf("JMP_IF_FALSE");
                for (int j = 0; j < indent; j++) printf("  ");
                printf(" R%d, %d, %d\n",
                       program.code[i], program.code[i+1], program.code[i+2]);
                i += 3;
                indent++;  // 增加缩进层级
                break;
            case OP_JMP:
                printf("JMP          %d, %d\n", program.code[i], program.code[i+1]);
                i += 2;
                break;
            case OP_NATIVE_CALL:
                printf("NATIVE_CALL  %d, %d, R%d\n",
                       program.code[i], program.code[i+1], program.code[i+2]);
                i += 3;
                break;
            case OP_HALT:
                printf("HALT\n");
                break;
            case OP_END_IF:
                indent--;  // 减少缩进层级
                printf("END_IF");
                for (int j = 0; j < indent; j++) printf("  ");
                printf("\n");
                break;
            default:
                printf("UNKNOWN     0x%02X\n", opcode);
                i++;
                break;
        }
    }
}