#include "vm.h"
#include <iostream>
#include <format>
#include <cmath>
#include <algorithm>

namespace camellya {

VM::VM() : current_frame(nullptr) {
    register_builtin_functions();
}

bool VM::execute(std::shared_ptr<Chunk> chunk) {
    if (!chunk) {
        runtime_error("Cannot execute null chunk");
        return false;
    }
    
    frames.clear();
    frames.emplace_back(chunk, 0);
    current_frame = &frames.back();
    stack.clear();
    
    try {
        return run();
    } catch (const std::exception& e) {
        runtime_error(e.what());
        return false;
    }
}

bool VM::run() {
    while (true) {
        // Debug: print instruction
        #ifdef DEBUG_TRACE_EXECUTION
        std::cout << "Stack: ";
        for (const auto& value : stack) {
            std::cout << "[ " << value->to_string() << " ]";
        }
        std::cout << std::endl;
        current_frame->chunk->disassemble_instruction(current_frame->ip);
        #endif
        
        OpCode instruction = static_cast<OpCode>(read_byte());
        
        switch (instruction) {
            case OpCode::OP_CONSTANT: {
                ValuePtr constant = read_constant();
                push(constant);
                break;
            }
            
            case OpCode::OP_NIL:
                push(std::make_shared<NilValue>());
                break;
            
            case OpCode::OP_TRUE:
                push(std::make_shared<BoolValue>(true));
                break;
            
            case OpCode::OP_FALSE:
                push(std::make_shared<BoolValue>(false));
                break;
            
            case OpCode::OP_ADD:
            case OpCode::OP_SUBTRACT:
            case OpCode::OP_MULTIPLY:
            case OpCode::OP_DIVIDE:
            case OpCode::OP_MODULO:
            case OpCode::OP_GREATER:
            case OpCode::OP_GREATER_EQUAL:
            case OpCode::OP_LESS:
            case OpCode::OP_LESS_EQUAL:
                if (!binary_op(instruction)) {
                    return false;
                }
                break;
            
            case OpCode::OP_EQUAL: {
                ValuePtr b = pop();
                ValuePtr a = pop();
                push(std::make_shared<BoolValue>(values_equal(a, b)));
                break;
            }
            
            case OpCode::OP_NOT_EQUAL: {
                ValuePtr b = pop();
                ValuePtr a = pop();
                push(std::make_shared<BoolValue>(!values_equal(a, b)));
                break;
            }
            
            case OpCode::OP_NEGATE: {
                ValuePtr value = pop();
                if (value->type() != Type::NUMBER) {
                    runtime_error("Operand must be a number.");
                    return false;
                }
                double num = std::dynamic_pointer_cast<NumberValue>(value)->value;
                push(std::make_shared<NumberValue>(-num));
                break;
            }
            
            case OpCode::OP_NOT: {
                ValuePtr value = pop();
                push(std::make_shared<BoolValue>(is_falsey(value)));
                break;
            }
            
            case OpCode::OP_POP:
                pop();
                break;
            
            case OpCode::OP_POPN: {
                uint8_t count = read_byte();
                for (int i = 0; i < count; i++) {
                    pop();
                }
                break;
            }
            
            case OpCode::OP_DUP:
                push(peek(0));
                break;
            
            case OpCode::OP_GET_GLOBAL: {
                std::string name = read_string();
                auto it = globals.find(name);
                if (it == globals.end()) {
                    // Fallback: check if 'this' is at slot 0 and has the property
                    ValuePtr receiver = stack[current_frame->stack_offset];
                    if (receiver && receiver->type() == Type::INSTANCE) {
                        auto instance = std::dynamic_pointer_cast<InstanceValue>(receiver);
                        if (instance->fields.find(name) != instance->fields.end()) {
                            push(instance->get(name));
                            break;
                        }
                    }
                    runtime_error("Undefined variable '" + name + "'.");
                    return false;
                }
                push(it->second);
                break;
            }
            
            case OpCode::OP_SET_GLOBAL: {
                std::string name = read_string();
                auto it = globals.find(name);
                if (it == globals.end()) {
                    // Fallback: check if 'this' is at slot 0 and has the property
                    ValuePtr receiver = stack[current_frame->stack_offset];
                    if (receiver && receiver->type() == Type::INSTANCE) {
                        auto instance = std::dynamic_pointer_cast<InstanceValue>(receiver);
                        if (instance->fields.find(name) != instance->fields.end()) {
                            instance->set(name, peek(0));
                            break;
                        }
                    }
                    runtime_error("Undefined variable '" + name + "'.");
                    return false;
                }
                it->second = peek(0);
                break;
            }
            
            case OpCode::OP_DEFINE_GLOBAL: {
                std::string name = read_string();
                globals[name] = pop();
                break;
            }
            
            case OpCode::OP_GET_LOCAL: {
                uint8_t slot = read_byte();
                push(stack[current_frame->stack_offset + slot]);
                break;
            }
            
            case OpCode::OP_SET_LOCAL: {
                uint8_t slot = read_byte();
                stack[current_frame->stack_offset + slot] = peek(0);
                break;
            }
            
            case OpCode::OP_JUMP: {
                uint16_t offset = read_short();
                current_frame->ip += offset;
                break;
            }
            
            case OpCode::OP_JUMP_IF_FALSE: {
                uint16_t offset = read_short();
                if (is_falsey(peek(0))) {
                    current_frame->ip += offset;
                }
                break;
            }
            
            case OpCode::OP_LOOP: {
                uint16_t offset = read_short();
                current_frame->ip -= offset;
                break;
            }
            
            case OpCode::OP_CALL: {
                uint8_t arg_count = read_byte();
                ValuePtr callee = peek(arg_count);
                
                if (callee->type() == Type::FUNCTION) {
                    auto func = std::dynamic_pointer_cast<FunctionValue>(callee);
                    
                    if (func->is_native) {
                        // Native function call
                        std::vector<ValuePtr> args;
                        if (func->bound_instance) {
                            args.push_back(func->bound_instance);
                        }
                        for (int i = arg_count - 1; i >= 0; i--) {
                            args.push_back(peek(i));
                        }
                        
                        ValuePtr result = func->native_func(args);
                        
                        // Pop arguments and function
                        for (int i = 0; i <= arg_count; i++) {
                            pop();
                        }
                        
                        push(result);
                    } else {
                        // Script function - call bytecode
                        if (frames.size() >= 64) {
                            runtime_error("Stack overflow.");
                            return false;
                        }
                        
                        // Parameters are already on the stack
                        // Function is at stack index: stack.size() - arg_count - 1
                        size_t stack_offset = stack.size() - arg_count - 1;
                        
                        // If it's a bound method, replace function on stack with 'this'
                        if (func->bound_instance) {
                            stack[stack_offset] = func->bound_instance;
                        }
                        
                        frames.emplace_back(func->chunk, stack_offset);
                        current_frame = &frames.back();
                    }
                } else if (callee->type() == Type::CLASS) {
                    // Class instantiation
                    auto klass = std::dynamic_pointer_cast<ClassValue>(callee);
                    auto instance = std::make_shared<InstanceValue>(klass);
                    
                    // The class is at stack[stack.size() - arg_count - 1]
                    // Replace it with the instance
                    stack[stack.size() - arg_count - 1] = instance;
                    
                    // Check if class has init method
                    auto init_method = klass->methods.find("init");
                    if (init_method != klass->methods.end()) {
                        auto init_func = init_method->second;
                        if (init_func->is_native) {
                            // ... native init (rare)
                            std::vector<ValuePtr> args;
                            for (int i = arg_count - 1; i >= 0; i--) {
                                args.push_back(peek(i));
                            }
                            init_func->native_func(args);
                            // pop arguments, instance stays on stack
                            for (int i = 0; i < arg_count; i++) pop();
                        } else {
                            // Script init - push a new frame
                            if (frames.size() >= 64) {
                                runtime_error("Stack overflow.");
                                return false;
                            }
                            
                            // stack_offset points to the instance we just put there
                            size_t stack_offset = stack.size() - arg_count - 1;
                            frames.emplace_back(init_func->chunk, stack_offset);
                            current_frame = &frames.back();
                        }
                    } else {
                        // No init, just pop arguments
                        for (int i = 0; i < arg_count; i++) {
                            pop();
                        }
                    }
                } else {
                    runtime_error("Can only call functions and classes.");
                    return false;
                }
                break;
            }
            
            case OpCode::OP_RETURN: {
                ValuePtr result = pop();
                
                // Pop the frame
                size_t stack_offset = current_frame->stack_offset;
                frames.pop_back();
                
                // Pop locals and function from the stack
                while (stack.size() > stack_offset) {
                    pop();
                }
                
                if (frames.empty()) {
                    return true;
                }
                
                current_frame = &frames.back();
                push(result);
                break;
            }
            
            case OpCode::OP_BUILD_LIST: {
                uint8_t count = read_byte();
                auto list = std::make_shared<ListValue>();
                
                std::vector<ValuePtr> elements;
                for (int i = 0; i < count; i++) {
                    elements.push_back(pop());
                }
                
                // Reverse because we popped in reverse order
                for (int i = count - 1; i >= 0; i--) {
                    list->push(elements[i]);
                }
                
                push(list);
                break;
            }
            
            case OpCode::OP_BUILD_MAP: {
                uint8_t count = read_byte();
                auto map = std::make_shared<MapValue>();
                
                for (int i = 0; i < count; i++) {
                    ValuePtr value = pop();
                    ValuePtr key = pop();
                    
                    if (key->type() != Type::STRING) {
                        runtime_error("Map keys must be strings.");
                        return false;
                    }
                    
                    std::string key_str = std::dynamic_pointer_cast<StringValue>(key)->value;
                    map->set(key_str, value);
                }
                
                push(map);
                break;
            }
            
            case OpCode::OP_INDEX: {
                ValuePtr index = pop();
                ValuePtr object = pop();
                
                if (object->type() == Type::LIST) {
                    auto list = std::dynamic_pointer_cast<ListValue>(object);
                    if (index->type() != Type::NUMBER) {
                        runtime_error("List index must be a number.");
                        return false;
                    }
                    size_t idx = static_cast<size_t>(
                        std::dynamic_pointer_cast<NumberValue>(index)->value);
                    push(list->get(idx));
                } else if (object->type() == Type::MAP) {
                    auto map = std::dynamic_pointer_cast<MapValue>(object);
                    if (index->type() != Type::STRING) {
                        runtime_error("Map key must be a string.");
                        return false;
                    }
                    std::string key = std::dynamic_pointer_cast<StringValue>(index)->value;
                    push(map->get(key));
                } else {
                    runtime_error("Only lists and maps support indexing.");
                    return false;
                }
                break;
            }
            
            case OpCode::OP_INDEX_SET: {
                ValuePtr value = pop();
                ValuePtr index = pop();
                ValuePtr object = pop();
                
                if (object->type() == Type::LIST) {
                    auto list = std::dynamic_pointer_cast<ListValue>(object);
                    if (index->type() != Type::NUMBER) {
                        runtime_error("List index must be a number.");
                        return false;
                    }
                    size_t idx = static_cast<size_t>(
                        std::dynamic_pointer_cast<NumberValue>(index)->value);
                    list->set(idx, value);
                    push(value);
                } else if (object->type() == Type::MAP) {
                    auto map = std::dynamic_pointer_cast<MapValue>(object);
                    if (index->type() != Type::STRING) {
                        runtime_error("Map key must be a string.");
                        return false;
                    }
                    std::string key = std::dynamic_pointer_cast<StringValue>(index)->value;
                    map->set(key, value);
                    push(value);
                } else {
                    runtime_error("Only lists and maps support index assignment.");
                    return false;
                }
                break;
            }
            
            case OpCode::OP_CLASS: {
                std::string name = read_string();
                auto klass = std::make_shared<ClassValue>(name);
                push(klass);
                break;
            }
            
            case OpCode::OP_GET_PROPERTY: {
                std::string name = read_string();
                ValuePtr object = pop();
                
                if (object->type() == Type::INSTANCE) {
                    auto instance = std::dynamic_pointer_cast<InstanceValue>(object);
                    push(instance->get(name));
                } else {
                    auto it = builtin_classes.find(object->type_name());
                    if (it != builtin_classes.end()) {
                        auto klass = it->second;
                        auto method_it = klass->methods.find(name);
                        if (method_it != klass->methods.end()) {
                            auto bound = std::make_shared<FunctionValue>(*method_it->second);
                            bound->bound_instance = object;
                            push(bound);
                            break;
                        }
                    }
                    runtime_error("Only instances and types with registered built-in classes have properties/methods.");
                    return false;
                }
                break;
            }
            
            case OpCode::OP_SET_PROPERTY: {
                std::string name = read_string();
                ValuePtr value = pop();
                ValuePtr object = pop();
                
                if (object->type() == Type::INSTANCE) {
                    auto instance = std::dynamic_pointer_cast<InstanceValue>(object);
                    instance->set(name, value);
                    push(value);
                } else {
                    runtime_error("Only instances have fields.");
                    return false;
                }
                break;
            }
            
            case OpCode::OP_PRINT: {
                std::cout << pop()->to_string() << std::endl;
                break;
            }
            
            case OpCode::OP_HALT:
                return true;
            
            default:
                runtime_error("Unknown opcode: " + 
                             std::to_string(static_cast<int>(instruction)));
                return false;
        }
    }
}

uint8_t VM::read_byte() {
    return current_frame->chunk->get_code(current_frame->ip++);
}

uint16_t VM::read_short() {
    uint16_t high = read_byte();
    uint16_t low = read_byte();
    return (high << 8) | low;
}

ValuePtr VM::read_constant() {
    uint8_t constant_idx = read_byte();
    return current_frame->chunk->get_constant(constant_idx);
}

std::string VM::read_string() {
    ValuePtr constant = read_constant();
    if (constant->type() != Type::STRING) {
        throw RuntimeError("Expected string constant");
    }
    return std::dynamic_pointer_cast<StringValue>(constant)->value;
}

void VM::runtime_error(const std::string& message) {
    error_message = message;
    std::cerr << "Runtime Error: " << message << std::endl;
    
    // Print stack trace
    for (auto it = frames.rbegin(); it != frames.rend(); ++it) {
        size_t instruction = it->ip - 1;
        int line = it->chunk->get_line(instruction);
        std::cerr << "[line " << line << "] in script" << std::endl;
    }
}

bool VM::is_falsey(ValuePtr value) {
    if (!value || value->type() == Type::NIL) return true;
    if (value->type() == Type::BOOL) {
        return !std::dynamic_pointer_cast<BoolValue>(value)->value;
    }
    return false;
}

bool VM::binary_op(OpCode op) {
    ValuePtr b = pop();
    ValuePtr a = pop();
    
    if (op == OpCode::OP_ADD) {
        if (a->type() == Type::NUMBER && b->type() == Type::NUMBER) {
            double av = std::dynamic_pointer_cast<NumberValue>(a)->value;
            double bv = std::dynamic_pointer_cast<NumberValue>(b)->value;
            push(std::make_shared<NumberValue>(av + bv));
            return true;
        }
        if (a->type() == Type::STRING && b->type() == Type::STRING) {
            std::string av = std::dynamic_pointer_cast<StringValue>(a)->value;
            std::string bv = std::dynamic_pointer_cast<StringValue>(b)->value;
            push(std::make_shared<StringValue>(av + bv));
            return true;
        }
        runtime_error("Operands must be two numbers or two strings.");
        return false;
    }
    
    // Other arithmetic operations require numbers
    if (a->type() != Type::NUMBER || b->type() != Type::NUMBER) {
        runtime_error("Operands must be numbers.");
        return false;
    }
    
    double av = std::dynamic_pointer_cast<NumberValue>(a)->value;
    double bv = std::dynamic_pointer_cast<NumberValue>(b)->value;
    
    switch (op) {
        case OpCode::OP_SUBTRACT:
            push(std::make_shared<NumberValue>(av - bv));
            break;
        case OpCode::OP_MULTIPLY:
            push(std::make_shared<NumberValue>(av * bv));
            break;
        case OpCode::OP_DIVIDE:
            if (bv == 0) {
                runtime_error("Division by zero.");
                return false;
            }
            push(std::make_shared<NumberValue>(av / bv));
            break;
        case OpCode::OP_MODULO:
            push(std::make_shared<NumberValue>(std::fmod(av, bv)));
            break;
        case OpCode::OP_GREATER:
            push(std::make_shared<BoolValue>(av > bv));
            break;
        case OpCode::OP_GREATER_EQUAL:
            push(std::make_shared<BoolValue>(av >= bv));
            break;
        case OpCode::OP_LESS:
            push(std::make_shared<BoolValue>(av < bv));
            break;
        case OpCode::OP_LESS_EQUAL:
            push(std::make_shared<BoolValue>(av <= bv));
            break;
        default:
            runtime_error("Unknown binary operator.");
            return false;
    }
    
    return true;
}

void VM::register_builtin_functions() {
    // print function
    auto print_func = std::make_shared<FunctionValue>("print", 
        [](const std::vector<ValuePtr>& args) -> ValuePtr {
            if (args.empty()) {
                std::cout << std::endl;
                return std::make_shared<NilValue>();
            }
            
            for (size_t i = 0; i < args.size(); ++i) {
                if (i > 0) std::cout << " ";
                std::cout << args[i]->to_string();
            }
            std::cout << std::endl;
            
            return std::make_shared<NilValue>();
        });
    globals["print"] = print_func;
    
    // len function
    auto len_func = std::make_shared<FunctionValue>("len",
        [](const std::vector<ValuePtr>& args) -> ValuePtr {
            if (args.size() != 1) {
                throw RuntimeError("len() expects 1 argument.");
            }
            
            auto& arg = args[0];
            if (arg->type() == Type::LIST) {
                auto list = std::dynamic_pointer_cast<ListValue>(arg);
                return std::make_shared<NumberValue>(static_cast<double>(list->size()));
            } else if (arg->type() == Type::STRING) {
                auto str = std::dynamic_pointer_cast<StringValue>(arg);
                return std::make_shared<NumberValue>(static_cast<double>(str->value.length()));
            } else if (arg->type() == Type::MAP) {
                auto map = std::dynamic_pointer_cast<MapValue>(arg);
                return std::make_shared<NumberValue>(static_cast<double>(map->pairs.size()));
            }
            
            throw RuntimeError("len() expects list, string, or map.");
        });
    globals["len"] = len_func;

    // Define list class
    auto list_klass = std::make_shared<ClassValue>("list");
    list_klass->add_method("push", std::make_shared<FunctionValue>("push",
        [](const std::vector<ValuePtr>& args) -> ValuePtr {
            if (args.size() != 2) throw RuntimeError("list.push() expects 1 argument.");
            auto list = std::dynamic_pointer_cast<ListValue>(args[0]);
            list->push(args[1]);
            return args[1];
        }));
    list_klass->add_method("pop", std::make_shared<FunctionValue>("pop",
        [](const std::vector<ValuePtr>& args) -> ValuePtr {
            if (args.size() != 1) throw RuntimeError("list.pop() expects 0 arguments.");
            auto list = std::dynamic_pointer_cast<ListValue>(args[0]);
            if (list->elements.empty()) return std::make_shared<NilValue>();
            ValuePtr val = list->elements.back();
            list->elements.pop_back();
            return val;
        }));
    list_klass->add_method("len", std::make_shared<FunctionValue>("len",
        [](const std::vector<ValuePtr>& args) -> ValuePtr {
            if (args.size() != 1) throw RuntimeError("list.len() expects 0 arguments.");
            auto list = std::dynamic_pointer_cast<ListValue>(args[0]);
            return std::make_shared<NumberValue>(static_cast<double>(list->size()));
        }));
    register_builtin_class("list", list_klass);

    // Define map class
    auto map_klass = std::make_shared<ClassValue>("map");
    map_klass->add_method("set", std::make_shared<FunctionValue>("set",
        [](const std::vector<ValuePtr>& args) -> ValuePtr {
            if (args.size() != 3) throw RuntimeError("map.set() expects 2 arguments.");
            auto map = std::dynamic_pointer_cast<MapValue>(args[0]);
            if (args[1]->type() != Type::STRING) throw RuntimeError("Map key must be a string.");
            std::string key = std::dynamic_pointer_cast<StringValue>(args[1])->value;
            map->set(key, args[2]);
            return args[2];
        }));
    map_klass->add_method("get", std::make_shared<FunctionValue>("get",
        [](const std::vector<ValuePtr>& args) -> ValuePtr {
            if (args.size() != 2) throw RuntimeError("map.get() expects 1 argument.");
            auto map = std::dynamic_pointer_cast<MapValue>(args[0]);
            if (args[1]->type() != Type::STRING) throw RuntimeError("Map key must be a string.");
            std::string key = std::dynamic_pointer_cast<StringValue>(args[1])->value;
            return map->get(key);
        }));
    map_klass->add_method("has", std::make_shared<FunctionValue>("has",
        [](const std::vector<ValuePtr>& args) -> ValuePtr {
            if (args.size() != 2) throw RuntimeError("map.has() expects 1 argument.");
            auto map = std::dynamic_pointer_cast<MapValue>(args[0]);
            if (args[1]->type() != Type::STRING) throw RuntimeError("Map key must be a string.");
            std::string key = std::dynamic_pointer_cast<StringValue>(args[1])->value;
            return std::make_shared<BoolValue>(map->has(key));
        }));
    map_klass->add_method("len", std::make_shared<FunctionValue>("len",
        [](const std::vector<ValuePtr>& args) -> ValuePtr {
            if (args.size() != 1) throw RuntimeError("map.len() expects 0 arguments.");
            auto map = std::dynamic_pointer_cast<MapValue>(args[0]);
            return std::make_shared<NumberValue>(static_cast<double>(map->pairs.size()));
        }));
    register_builtin_class("map", map_klass);
}

void VM::register_native_function(const std::string& name, NativeFunction func) {
    auto func_value = std::make_shared<FunctionValue>(name, func);
    globals[name] = func_value;
}

void VM::register_builtin_class(const std::string& type_name, std::shared_ptr<ClassValue> klass) {
    builtin_classes[type_name] = klass;
    globals[klass->name] = klass;
}

void VM::set_global(const std::string& name, ValuePtr value) {
    globals[name] = value;
}

ValuePtr VM::get_global(const std::string& name) {
    auto it = globals.find(name);
    if (it == globals.end()) {
        return nullptr;
    }
    return it->second;
}

void VM::push(ValuePtr value) {
    stack.push_back(value);
}

ValuePtr VM::pop() {
    if (stack.empty()) {
        throw RuntimeError("Stack underflow");
    }
    ValuePtr value = stack.back();
    stack.pop_back();
    return value;
}

ValuePtr VM::peek(int distance) {
    if (distance >= static_cast<int>(stack.size())) {
        throw RuntimeError("Stack underflow in peek");
    }
    return stack[stack.size() - 1 - distance];
}

} // namespace camellya