#include "interpreter.h"
#include <iostream>
#include <format>
#include <cmath>

namespace camellya {

Interpreter::Interpreter() {
    global_environment = std::make_shared<Environment>();
    environment = global_environment;
    register_native_functions();
}

void Interpreter::register_native_functions() {
    // print function - supports format strings
    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>();
            }
            
            // Simple print: just concatenate all arguments
            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>();
        });
    global_environment->define("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.");
        });
    global_environment->define("len", len_func);
}

void Interpreter::execute(const Program& program) {
    for (const auto& stmt : program.statements) {
        execute_statement(*stmt);
    }
}

ValuePtr Interpreter::evaluate(const Expr& expr) {
    if (auto* binary = dynamic_cast<const BinaryExpr*>(&expr)) {
        return eval_binary(*binary);
    } else if (auto* unary = dynamic_cast<const UnaryExpr*>(&expr)) {
        return eval_unary(*unary);
    } else if (auto* literal = dynamic_cast<const LiteralExpr*>(&expr)) {
        return eval_literal(*literal);
    } else if (auto* variable = dynamic_cast<const VariableExpr*>(&expr)) {
        return eval_variable(*variable);
    } else if (auto* assign = dynamic_cast<const AssignExpr*>(&expr)) {
        return eval_assign(*assign);
    } else if (auto* call = dynamic_cast<const CallExpr*>(&expr)) {
        return eval_call(*call);
    } else if (auto* get = dynamic_cast<const GetExpr*>(&expr)) {
        return eval_get(*get);
    } else if (auto* set = dynamic_cast<const SetExpr*>(&expr)) {
        return eval_set(*set);
    } else if (auto* index = dynamic_cast<const IndexExpr*>(&expr)) {
        return eval_index(*index);
    } else if (auto* index_set = dynamic_cast<const IndexSetExpr*>(&expr)) {
        return eval_index_set(*index_set);
    } else if (auto* list = dynamic_cast<const ListExpr*>(&expr)) {
        return eval_list(*list);
    } else if (auto* map = dynamic_cast<const MapExpr*>(&expr)) {
        return eval_map(*map);
    }
    
    throw RuntimeError("Unknown expression type.");
}

void Interpreter::execute_statement(const Stmt& stmt) {
    if (auto* expr_stmt = dynamic_cast<const ExprStmt*>(&stmt)) {
        exec_expr_stmt(*expr_stmt);
    } else if (auto* var_decl = dynamic_cast<const VarDecl*>(&stmt)) {
        exec_var_decl(*var_decl);
    } else if (auto* block = dynamic_cast<const BlockStmt*>(&stmt)) {
        exec_block(*block);
    } else if (auto* if_stmt = dynamic_cast<const IfStmt*>(&stmt)) {
        exec_if(*if_stmt);
    } else if (auto* while_stmt = dynamic_cast<const WhileStmt*>(&stmt)) {
        exec_while(*while_stmt);
    } else if (auto* for_stmt = dynamic_cast<const ForStmt*>(&stmt)) {
        exec_for(*for_stmt);
    } else if (auto* return_stmt = dynamic_cast<const ReturnStmt*>(&stmt)) {
        exec_return(*return_stmt);
    } else if (auto* func_decl = dynamic_cast<const FunctionDecl*>(&stmt)) {
        exec_function_decl(*func_decl);
    } else if (auto* class_decl = dynamic_cast<const ClassDecl*>(&stmt)) {
        exec_class_decl(*class_decl);
    }
}

ValuePtr Interpreter::eval_binary(const BinaryExpr& expr) {
    ValuePtr left = evaluate(*expr.left);
    ValuePtr right = evaluate(*expr.right);
    
    if (expr.op == "+") {
        if (left->type() == Type::NUMBER && right->type() == Type::NUMBER) {
            double l = std::dynamic_pointer_cast<NumberValue>(left)->value;
            double r = std::dynamic_pointer_cast<NumberValue>(right)->value;
            return std::make_shared<NumberValue>(l + r);
        }
        
        if (left->type() == Type::STRING && right->type() == Type::STRING) {
            const auto& l = std::dynamic_pointer_cast<StringValue>(left)->value;
            const auto& r = std::dynamic_pointer_cast<StringValue>(right)->value;
            return std::make_shared<StringValue>(l + r);
        }
        
        throw RuntimeError("Operands of '+' must be both numbers or both strings.");
    }
    
    if (expr.op == "-" || expr.op == "*" || 
        expr.op == "/" || expr.op == "%") {
        if (left->type() != Type::NUMBER || right->type() != Type::NUMBER) {
            throw RuntimeError("Operands must be numbers.");
        }
        
        double l = std::dynamic_pointer_cast<NumberValue>(left)->value;
        double r = std::dynamic_pointer_cast<NumberValue>(right)->value;
        
        if (expr.op == "-") return std::make_shared<NumberValue>(l - r);
        if (expr.op == "*") return std::make_shared<NumberValue>(l * r);
        if (expr.op == "/") {
            if (r == 0) throw RuntimeError("Division by zero.");
            return std::make_shared<NumberValue>(l / r);
        }
        if (expr.op == "%") return std::make_shared<NumberValue>(std::fmod(l, r));
    }
    
    if (expr.op == "==" ) {
        return std::make_shared<BoolValue>(values_equal(left, right));
    }
    if (expr.op == "!=") {
        return std::make_shared<BoolValue>(!values_equal(left, right));
    }
    
    if (expr.op == "<" || expr.op == "<=" || expr.op == ">" || expr.op == ">=") {
        if (left->type() != Type::NUMBER || right->type() != Type::NUMBER) {
            throw RuntimeError("Operands must be numbers.");
        }
        
        double l = std::dynamic_pointer_cast<NumberValue>(left)->value;
        double r = std::dynamic_pointer_cast<NumberValue>(right)->value;
        
        if (expr.op == "<") return std::make_shared<BoolValue>(l < r);
        if (expr.op == "<=") return std::make_shared<BoolValue>(l <= r);
        if (expr.op == ">") return std::make_shared<BoolValue>(l > r);
        if (expr.op == ">=") return std::make_shared<BoolValue>(l >= r);
    }
    
    if (expr.op == "and") {
        if (!is_truthy(left)) return left;
        return right;
    }
    
    if (expr.op == "or") {
        if (is_truthy(left)) return left;
        return right;
    }
    
    throw RuntimeError("Unknown binary operator: " + expr.op);
}

ValuePtr Interpreter::eval_unary(const UnaryExpr& expr) {
    ValuePtr operand = evaluate(*expr.operand);
    
    if (expr.op == "-") {
        if (operand->type() != Type::NUMBER) {
            throw RuntimeError("Operand must be a number.");
        }
        double value = std::dynamic_pointer_cast<NumberValue>(operand)->value;
        return std::make_shared<NumberValue>(-value);
    }
    
    if (expr.op == "!") {
        return std::make_shared<BoolValue>(!is_truthy(operand));
    }
    
    throw RuntimeError("Unknown unary operator: " + expr.op);
}

ValuePtr Interpreter::eval_literal(const LiteralExpr& expr) {
    return std::visit([](auto&& arg) -> ValuePtr {
        using T = std::decay_t<decltype(arg)>;
        if constexpr (std::is_same_v<T, double>) {
            return std::make_shared<NumberValue>(arg);
        } else if constexpr (std::is_same_v<T, std::string>) {
            return std::make_shared<StringValue>(arg);
        } else if constexpr (std::is_same_v<T, bool>) {
            return std::make_shared<BoolValue>(arg);
        } else {
            return std::make_shared<NilValue>();
        }
    }, expr.value);
}

ValuePtr Interpreter::eval_variable(const VariableExpr& expr) {
    return environment->get(expr.name);
}

ValuePtr Interpreter::eval_assign(const AssignExpr& expr) {
    ValuePtr value = evaluate(*expr.value);
    environment->set(expr.name, value);
    return value;
}

ValuePtr Interpreter::eval_call(const CallExpr& expr) {
    ValuePtr callee = evaluate(*expr.callee);
    
    std::vector<ValuePtr> arguments;
    for (const auto& arg : expr.arguments) {
        arguments.push_back(evaluate(*arg));
    }
    
    if (callee->type() == Type::FUNCTION) {
        auto func = std::dynamic_pointer_cast<FunctionValue>(callee);
        return call_function(*func, arguments);
    } else if (callee->type() == Type::CLASS) {
        // Class instantiation
        auto klass = std::dynamic_pointer_cast<ClassValue>(callee);
        auto instance = std::make_shared<InstanceValue>(klass);
        
        // If there is an init method, call it like a constructor
        ValuePtr init_val = instance->get("init");
        if (init_val && init_val->type() == Type::FUNCTION) {
            auto init_func = std::dynamic_pointer_cast<FunctionValue>(init_val);
            call_function(*init_func, arguments);
        }
        
        return instance;
    }
    
    throw RuntimeError("Can only call functions and classes.");
}

ValuePtr Interpreter::eval_get(const GetExpr& expr) {
    ValuePtr object = evaluate(*expr.object);
    
    if (object->type() == Type::INSTANCE) {
        auto instance = std::dynamic_pointer_cast<InstanceValue>(object);
        return instance->get(expr.name);
    }
    
    throw RuntimeError("Only instances have properties.");
}

ValuePtr Interpreter::eval_set(const SetExpr& expr) {
    ValuePtr object = evaluate(*expr.object);
    
    if (object->type() == Type::INSTANCE) {
        auto instance = std::dynamic_pointer_cast<InstanceValue>(object);
        ValuePtr value = evaluate(*expr.value);
        instance->set(expr.name, value);
        return value;
    }
    
    throw RuntimeError("Only instances have fields.");
}

ValuePtr Interpreter::eval_index(const IndexExpr& expr) {
    ValuePtr object = evaluate(*expr.object);
    ValuePtr index = evaluate(*expr.index);
    
    if (object->type() == Type::LIST) {
        auto list = std::dynamic_pointer_cast<ListValue>(object);
        if (index->type() != Type::NUMBER) {
            throw RuntimeError("List index must be a number.");
        }
        size_t idx = static_cast<size_t>(std::dynamic_pointer_cast<NumberValue>(index)->value);
        return list->get(idx);
    } else if (object->type() == Type::MAP) {
        auto map = std::dynamic_pointer_cast<MapValue>(object);
        if (index->type() != Type::STRING) {
            throw RuntimeError("Map key must be a string.");
        }
        std::string key = std::dynamic_pointer_cast<StringValue>(index)->value;
        return map->get(key);
    }
    
    throw RuntimeError("Only lists and maps support indexing.");
}

ValuePtr Interpreter::eval_index_set(const IndexSetExpr& expr) {
    ValuePtr object = evaluate(*expr.object);
    ValuePtr index = evaluate(*expr.index);
    ValuePtr value = evaluate(*expr.value);
    
    if (object->type() == Type::LIST) {
        auto list = std::dynamic_pointer_cast<ListValue>(object);
        if (index->type() != Type::NUMBER) {
            throw RuntimeError("List index must be a number.");
        }
        size_t idx = static_cast<size_t>(std::dynamic_pointer_cast<NumberValue>(index)->value);
        list->set(idx, value);
        return value;
    } else if (object->type() == Type::MAP) {
        auto map = std::dynamic_pointer_cast<MapValue>(object);
        if (index->type() != Type::STRING) {
            throw RuntimeError("Map key must be a string.");
        }
        std::string key = std::dynamic_pointer_cast<StringValue>(index)->value;
        map->set(key, value);
        return value;
    }
    
    throw RuntimeError("Only lists and maps support index assignment.");
}

ValuePtr Interpreter::eval_list(const ListExpr& expr) {
    auto list = std::make_shared<ListValue>();
    for (const auto& elem : expr.elements) {
        list->push(evaluate(*elem));
    }
    return list;
}

ValuePtr Interpreter::eval_map(const MapExpr& expr) {
    auto map = std::make_shared<MapValue>();
    for (const auto& [key_expr, value_expr] : expr.pairs) {
        ValuePtr key = evaluate(*key_expr);
        ValuePtr value = evaluate(*value_expr);
        
        if (key->type() != Type::STRING) {
            throw RuntimeError("Map keys must be strings.");
        }
        
        std::string key_str = std::dynamic_pointer_cast<StringValue>(key)->value;
        map->set(key_str, value);
    }
    return map;
}

void Interpreter::exec_expr_stmt(const ExprStmt& stmt) {
    evaluate(*stmt.expression);
}

void Interpreter::exec_var_decl(const VarDecl& stmt) {
    ValuePtr value;
    
    if (stmt.initializer) {
        value = evaluate(*stmt.initializer);
    } else {
        value = create_default_value(stmt.type_name);
    }
    
    environment->define(stmt.name, value);
}

void Interpreter::exec_block(const BlockStmt& stmt) {
    auto previous = environment;
    environment = std::make_shared<Environment>(environment);
    
    try {
        for (const auto& statement : stmt.statements) {
            execute_statement(*statement);
        }
    } catch (...) {
        environment = previous;
        throw;
    }
    
    environment = previous;
}

void Interpreter::exec_if(const IfStmt& stmt) {
    ValuePtr condition = evaluate(*stmt.condition);
    
    if (is_truthy(condition)) {
        execute_statement(*stmt.then_branch);
    } else if (stmt.else_branch) {
        execute_statement(*stmt.else_branch);
    }
}

void Interpreter::exec_while(const WhileStmt& stmt) {
    while (is_truthy(evaluate(*stmt.condition))) {
        execute_statement(*stmt.body);
    }
}

void Interpreter::exec_for(const ForStmt& stmt) {
    auto previous = environment;
    environment = std::make_shared<Environment>(environment);
    
    try {
        if (stmt.initializer) {
            execute_statement(*stmt.initializer);
        }
        
        while (!stmt.condition || is_truthy(evaluate(*stmt.condition))) {
            execute_statement(*stmt.body);
            
            if (stmt.increment) {
                evaluate(*stmt.increment);
            }
        }
    } catch (...) {
        environment = previous;
        throw;
    }
    
    environment = previous;
}

void Interpreter::exec_return(const ReturnStmt& stmt) {
    ValuePtr value = stmt.value ? evaluate(*stmt.value) : std::make_shared<NilValue>();
    throw ReturnException(value);
}

void Interpreter::exec_function_decl(const FunctionDecl& stmt) {
    auto func_decl = std::make_shared<FunctionDecl>(stmt);
    auto func = std::make_shared<FunctionValue>(stmt.name, func_decl);
    environment->define(stmt.name, func);
}

void Interpreter::exec_class_decl(const ClassDecl& stmt) {
    auto klass = std::make_shared<ClassValue>(stmt.name);
    
    for (const auto& member : stmt.members) {
        if (auto* var_decl = dynamic_cast<VarDecl*>(member.get())) {
            klass->add_field(var_decl->name, var_decl->type_name);
        } else if (auto* func_decl = dynamic_cast<FunctionDecl*>(member.get())) {
            auto func_decl_ptr = std::make_shared<FunctionDecl>(*func_decl);
            auto func = std::make_shared<FunctionValue>(func_decl->name, func_decl_ptr);
            klass->add_method(func_decl->name, func);
        }
    }
    
    environment->define(stmt.name, klass);
}

ValuePtr Interpreter::call_function(const FunctionValue& func, const std::vector<ValuePtr>& arguments) {
    if (func.is_native) {
        return func.native_func(arguments);
    }
    
    // Bind parameters
    if (func.declaration->parameters.size() != arguments.size()) {
        throw RuntimeError(std::format("Expected {} arguments but got {}.", 
                                      func.declaration->parameters.size(), 
                                      arguments.size()));
    }
    
    auto previous = environment;
    environment = std::make_shared<Environment>(global_environment);
    
    // Bind 'this' for methods
    if (func.bound_instance) {
        environment->define("this", func.bound_instance);
    }
    
    // Bind parameters to the new environment
    for (size_t i = 0; i < arguments.size(); ++i) {
        environment->define(func.declaration->parameters[i].second, arguments[i]);
    }
    
    try {
        execute_statement(*func.declaration->body);
    } catch (const ReturnException& ret) {
        environment = previous;
        return ret.value;
    }
    
    environment = previous;
    return std::make_shared<NilValue>();
}

ValuePtr Interpreter::create_default_value(const std::string& type_name) {
    if (type_name == "number") {
        return std::make_shared<NumberValue>(0.0);
    } else if (type_name == "string") {
        return std::make_shared<StringValue>("");
    } else if (type_name == "bool") {
        return std::make_shared<BoolValue>(false);
    } else if (type_name == "list") {
        return std::make_shared<ListValue>();
    } else if (type_name == "map") {
        return std::make_shared<MapValue>();
    } else if (environment->has(type_name)) {
        // It's a class type
        ValuePtr klass = environment->get(type_name);
        if (klass->type() == Type::CLASS) {
            auto klass_ptr = std::dynamic_pointer_cast<ClassValue>(klass);
            auto instance = std::make_shared<InstanceValue>(klass_ptr);
            
            // If the class defines init(), call it with no arguments
            ValuePtr init_val = instance->get("init");
            if (init_val && init_val->type() == Type::FUNCTION) {
                auto init_func = std::dynamic_pointer_cast<FunctionValue>(init_val);
                call_function(*init_func, {});
            }
            
            return instance;
        }
    }
    
    return std::make_shared<NilValue>();
}

} // namespace camellya