#define _POSIX_C_SOURCE 200809L
#include "semantic.h"
#include "../utils/safe_alloc.h"
#include <stdlib.h>
#include <string.h>
#include <stdio.h>

static void _rava_semantic_error(RavaSemanticAnalyzer_t *analyzer, const char *message, int line, int col);
static bool _rava_semantic_analyze_node(RavaSemanticAnalyzer_t *analyzer, RavaASTNode_t *node);
static RavaType_t* _rava_semantic_check_expression(RavaSemanticAnalyzer_t *analyzer, RavaASTNode_t *expr);

RavaSemanticAnalyzer_t* rava_semantic_analyzer_create() {
    RavaSemanticAnalyzer_t *analyzer = calloc(1, sizeof(RavaSemanticAnalyzer_t));
    if (!analyzer) return NULL;
    analyzer->symbol_table = rava_symbol_table_create();
    if (!analyzer->symbol_table) {
        free(analyzer);
        return NULL;
    }
    analyzer->error_message = NULL;
    analyzer->had_error = false;
    analyzer->error_count = 0;
    return analyzer;
}

void rava_semantic_analyzer_destroy(RavaSemanticAnalyzer_t *analyzer) {
    if (!analyzer) return;
    rava_symbol_table_destroy(analyzer->symbol_table);
    free(analyzer->error_message);
    free(analyzer);
}

static void _rava_semantic_error(RavaSemanticAnalyzer_t *analyzer, const char *message, int line, int col) {
    analyzer->had_error = true;
    analyzer->error_count++;

    if (analyzer->error_message) free(analyzer->error_message);
    analyzer->error_message = malloc(RAVA_ERROR_BUFFER_SIZE);
    if (!analyzer->error_message) return;
    snprintf(analyzer->error_message, RAVA_ERROR_BUFFER_SIZE, "Semantic error at line %d, column %d: %s",
             line, col, message);
}

static RavaType_t* _rava_semantic_resolve_type(RavaSemanticAnalyzer_t *analyzer, RavaASTNode_t *type_node) {
    (void)analyzer;
    if (!type_node || type_node->type != RAVA_AST_TYPE) return NULL;

    RavaType_t *base_type = rava_type_from_name(type_node->data.type.type_name);

    if (type_node->data.type.is_array) {
        return rava_type_create_array(base_type, type_node->data.type.array_dimensions);
    }

    return base_type;
}

static bool _rava_semantic_analyze_class(RavaSemanticAnalyzer_t *analyzer, RavaASTNode_t *class_node) {
    if (class_node->type != RAVA_AST_CLASS_DECL) return false;

    char *class_name = class_node->data.class_decl.name;
    RavaType_t *class_type = rava_type_create_class(class_name);

    RavaSymbol_t *class_symbol = rava_symbol_create(RAVA_SYMBOL_CLASS, class_name, class_type);
    class_symbol->modifiers = class_node->data.class_decl.modifiers;
    class_symbol->modifiers_count = class_node->data.class_decl.modifiers_count;
    class_symbol->declaration = class_node;

    if (!rava_symbol_table_define(analyzer->symbol_table, class_symbol)) {
        _rava_semantic_error(analyzer, "Class already defined", class_node->line, class_node->column);
        rava_symbol_destroy(class_symbol);
        return false;
    }

    rava_symbol_table_enter_scope(analyzer->symbol_table, class_name);

    for (size_t i = 0; i < class_node->children_count; i++) {
        _rava_semantic_analyze_node(analyzer, class_node->children[i]);
    }

    rava_symbol_table_exit_scope(analyzer->symbol_table);

    return !analyzer->had_error;
}

static bool _rava_semantic_analyze_enum(RavaSemanticAnalyzer_t *analyzer, RavaASTNode_t *enum_node) {
    if (enum_node->type != RAVA_AST_ENUM_DECL) return false;

    char *enum_name = enum_node->data.enum_decl.name;
    RavaType_t *enum_type = rava_type_create_class(enum_name);

    RavaSymbol_t *enum_symbol = rava_symbol_create(RAVA_SYMBOL_ENUM, enum_name, enum_type);
    enum_symbol->modifiers = enum_node->data.enum_decl.modifiers;
    enum_symbol->modifiers_count = enum_node->data.enum_decl.modifiers_count;
    enum_symbol->declaration = enum_node;

    if (!rava_symbol_table_define(analyzer->symbol_table, enum_symbol)) {
        _rava_semantic_error(analyzer, "Enum already defined", enum_node->line, enum_node->column);
        rava_symbol_destroy(enum_symbol);
        return false;
    }

    rava_symbol_table_enter_scope(analyzer->symbol_table, enum_name);

    for (size_t i = 0; i < enum_node->data.enum_decl.constants_count; i++) {
        char *const_name = enum_node->data.enum_decl.constants[i];
        RavaType_t *const_type = rava_type_create_class(enum_name);
        RavaSymbol_t *const_symbol = rava_symbol_create(RAVA_SYMBOL_FIELD, const_name, const_type);
        rava_symbol_table_define(analyzer->symbol_table, const_symbol);
    }

    for (size_t i = 0; i < enum_node->children_count; i++) {
        _rava_semantic_analyze_node(analyzer, enum_node->children[i]);
    }

    rava_symbol_table_exit_scope(analyzer->symbol_table);

    return !analyzer->had_error;
}

static bool _rava_semantic_analyze_method(RavaSemanticAnalyzer_t *analyzer, RavaASTNode_t *method_node) {
    if (method_node->type != RAVA_AST_METHOD_DECL) return false;

    char *method_name = method_node->data.method_decl.name;
    RavaType_t *return_type = _rava_semantic_resolve_type(analyzer, method_node->data.method_decl.return_type);

    RavaSymbol_t *method_symbol = rava_symbol_create(RAVA_SYMBOL_METHOD, method_name, return_type);
    method_symbol->modifiers = method_node->data.method_decl.modifiers;
    method_symbol->modifiers_count = method_node->data.method_decl.modifiers_count;
    method_symbol->declaration = method_node;

    if (!rava_symbol_table_define(analyzer->symbol_table, method_symbol)) {
        char error_msg[256];
        snprintf(error_msg, sizeof(error_msg), "Method '%s' already defined", method_name);
        _rava_semantic_error(analyzer, error_msg, method_node->line, method_node->column);
        rava_symbol_destroy(method_symbol);
        return false;
    }

    rava_symbol_table_enter_scope(analyzer->symbol_table, method_name);

    for (size_t i = 0; i < method_node->children_count; i++) {
        RavaASTNode_t *child = method_node->children[i];

        if (child->type == RAVA_AST_PARAM_DECL) {
            char *param_name = child->data.var_decl.name;
            RavaType_t *param_type = _rava_semantic_resolve_type(analyzer, child->data.var_decl.type);

            RavaSymbol_t *param_symbol = rava_symbol_create(RAVA_SYMBOL_PARAMETER, param_name, param_type);
            param_symbol->declaration = child;

            if (!rava_symbol_table_define(analyzer->symbol_table, param_symbol)) {
                char error_msg[256];
                snprintf(error_msg, sizeof(error_msg), "Parameter '%s' already defined", param_name);
                _rava_semantic_error(analyzer, error_msg, child->line, child->column);
                rava_symbol_destroy(param_symbol);
            }
        } else if (child->type == RAVA_AST_BLOCK_STMT) {
            _rava_semantic_analyze_node(analyzer, child);
        }
    }

    rava_symbol_table_exit_scope(analyzer->symbol_table);

    return !analyzer->had_error;
}

static bool _rava_semantic_analyze_field(RavaSemanticAnalyzer_t *analyzer, RavaASTNode_t *field_node) {
    if (field_node->type != RAVA_AST_FIELD_DECL) return false;

    char *field_name = field_node->data.field_decl.name;
    RavaType_t *field_type = _rava_semantic_resolve_type(analyzer, field_node->data.field_decl.type);

    RavaSymbol_t *field_symbol = rava_symbol_create(RAVA_SYMBOL_FIELD, field_name, field_type);
    field_symbol->declaration = field_node;
    field_symbol->modifiers = field_node->data.field_decl.modifiers;
    field_symbol->modifiers_count = field_node->data.field_decl.modifiers_count;

    if (!rava_symbol_table_define(analyzer->symbol_table, field_symbol)) {
        char error_msg[256];
        snprintf(error_msg, sizeof(error_msg), "Field '%s' already defined", field_name);
        _rava_semantic_error(analyzer, error_msg, field_node->line, field_node->column);
        rava_symbol_destroy(field_symbol);
        return false;
    }

    if (field_node->data.field_decl.initializer) {
        RavaType_t *init_type = _rava_semantic_check_expression(analyzer, field_node->data.field_decl.initializer);
        if (init_type && !rava_type_is_assignable_to(init_type, field_type)) {
            char error_msg[256];
            snprintf(error_msg, sizeof(error_msg),
                    "Cannot assign '%s' to field of type '%s'",
                    rava_type_to_string(init_type),
                    rava_type_to_string(field_type));
            _rava_semantic_error(analyzer, error_msg, field_node->line, field_node->column);
        }
    }

    return !analyzer->had_error;
}

static bool _rava_semantic_analyze_var_decl(RavaSemanticAnalyzer_t *analyzer, RavaASTNode_t *var_node) {
    if (var_node->type != RAVA_AST_VAR_DECL) return false;

    char *var_name = var_node->data.var_decl.name;
    RavaType_t *var_type = _rava_semantic_resolve_type(analyzer, var_node->data.var_decl.type);

    RavaSymbol_t *var_symbol = rava_symbol_create(RAVA_SYMBOL_VARIABLE, var_name, var_type);
    var_symbol->declaration = var_node;

    if (!rava_symbol_table_define(analyzer->symbol_table, var_symbol)) {
        char error_msg[256];
        snprintf(error_msg, sizeof(error_msg), "Variable '%s' already defined", var_name);
        _rava_semantic_error(analyzer, error_msg, var_node->line, var_node->column);
        rava_symbol_destroy(var_symbol);
        return false;
    }

    if (var_node->data.var_decl.initializer) {
        RavaType_t *init_type = _rava_semantic_check_expression(analyzer, var_node->data.var_decl.initializer);
        if (init_type && !rava_type_is_assignable_to(init_type, var_type)) {
            char error_msg[256];
            snprintf(error_msg, sizeof(error_msg),
                    "Cannot assign '%s' to variable of type '%s'",
                    rava_type_to_string(init_type),
                    rava_type_to_string(var_type));
            _rava_semantic_error(analyzer, error_msg, var_node->line, var_node->column);
        }
    }

    return !analyzer->had_error;
}

static RavaType_t* _rava_semantic_check_expression(RavaSemanticAnalyzer_t *analyzer, RavaASTNode_t *expr) {
    if (!expr) return NULL;

    switch (expr->type) {
        case RAVA_AST_LITERAL_EXPR: {
            switch (expr->data.literal.literal_type) {
                case RAVA_TOKEN_LITERAL_INTEGER:
                    return rava_type_create_primitive(RAVA_TYPE_INT);
                case RAVA_TOKEN_LITERAL_LONG:
                    return rava_type_create_primitive(RAVA_TYPE_LONG);
                case RAVA_TOKEN_LITERAL_FLOAT:
                    return rava_type_create_primitive(RAVA_TYPE_FLOAT);
                case RAVA_TOKEN_LITERAL_DOUBLE:
                    return rava_type_create_primitive(RAVA_TYPE_DOUBLE);
                case RAVA_TOKEN_LITERAL_STRING:
                    return rava_type_create_class("String");
                case RAVA_TOKEN_LITERAL_CHARACTER:
                    return rava_type_create_primitive(RAVA_TYPE_CHAR);
                case RAVA_TOKEN_LITERAL_TRUE:
                case RAVA_TOKEN_LITERAL_FALSE:
                    return rava_type_create_primitive(RAVA_TYPE_BOOLEAN);
                case RAVA_TOKEN_LITERAL_NULL:
                    return rava_type_create_primitive(RAVA_TYPE_NULL);
                default:
                    return NULL;
            }
        }

        case RAVA_AST_IDENTIFIER_EXPR: {
            char *name = expr->data.identifier.name;

            if (strcmp(name, "System") == 0) {
                return rava_type_from_name("System");
            }
            if (strcmp(name, "Files") == 0) {
                return rava_type_from_name("Files");
            }

            RavaSymbol_t *symbol = rava_symbol_table_resolve(analyzer->symbol_table, name);
            if (!symbol) {
                char error_msg[256];
                snprintf(error_msg, sizeof(error_msg), "Undefined identifier '%s'", name);
                _rava_semantic_error(analyzer, error_msg, expr->line, expr->column);
                return NULL;
            }
            return symbol->type;
        }

        case RAVA_AST_BINARY_EXPR: {
            RavaType_t *left_type = _rava_semantic_check_expression(analyzer, expr->data.binary.left);
            RavaType_t *right_type = _rava_semantic_check_expression(analyzer, expr->data.binary.right);

            if (!left_type || !right_type) return NULL;

            if (expr->data.binary.op >= RAVA_BINOP_EQ && expr->data.binary.op <= RAVA_BINOP_GE) {
                return rava_type_create_primitive(RAVA_TYPE_BOOLEAN);
            }

            if (expr->data.binary.op == RAVA_BINOP_AND || expr->data.binary.op == RAVA_BINOP_OR) {
                if (left_type->kind != RAVA_TYPE_BOOLEAN || right_type->kind != RAVA_TYPE_BOOLEAN) {
                    _rava_semantic_error(analyzer, "Logical operators require boolean operands",
                                       expr->line, expr->column);
                    return NULL;
                }
                return rava_type_create_primitive(RAVA_TYPE_BOOLEAN);
            }

            return rava_type_binary_result(left_type, right_type);
        }

        case RAVA_AST_UNARY_EXPR: {
            return _rava_semantic_check_expression(analyzer, expr->data.unary.operand);
        }

        case RAVA_AST_ASSIGN_EXPR: {
            RavaType_t *target_type = _rava_semantic_check_expression(analyzer, expr->data.assign.target);
            RavaType_t *value_type = _rava_semantic_check_expression(analyzer, expr->data.assign.value);

            if (target_type && value_type && !rava_type_is_assignable_to(value_type, target_type)) {
                char error_msg[256];
                snprintf(error_msg, sizeof(error_msg),
                        "Cannot assign '%s' to '%s'",
                        rava_type_to_string(value_type),
                        rava_type_to_string(target_type));
                _rava_semantic_error(analyzer, error_msg, expr->line, expr->column);
            }

            return target_type;
        }

        case RAVA_AST_CALL_EXPR: {
            _rava_semantic_check_expression(analyzer, expr->data.call.callee);
            for (size_t i = 0; i < expr->data.call.arguments_count; i++) {
                _rava_semantic_check_expression(analyzer, expr->data.call.arguments[i]);
            }
            if (expr->data.call.callee->type == RAVA_AST_MEMBER_ACCESS_EXPR) {
                RavaASTNode_t *member = expr->data.call.callee;
                const char *method = member->data.member_access.member;
                if (member->data.member_access.object->type == RAVA_AST_IDENTIFIER_EXPR &&
                    strcmp(member->data.member_access.object->data.identifier.name, "Files") == 0) {
                    if (strcmp(method, "read") == 0) {
                        return rava_type_create_class("String");
                    } else if (strcmp(method, "write") == 0 || strcmp(method, "exists") == 0 ||
                               strcmp(method, "delete") == 0) {
                        return rava_type_create_primitive(RAVA_TYPE_BOOLEAN);
                    }
                }
                if (member->data.member_access.object->type == RAVA_AST_IDENTIFIER_EXPR &&
                    strcmp(member->data.member_access.object->data.identifier.name, "System") == 0) {
                    if (strcmp(method, "currentTimeMillis") == 0 || strcmp(method, "nanoTime") == 0) {
                        return rava_type_create_primitive(RAVA_TYPE_LONG);
                    }
                }
                if (member->data.member_access.object->type == RAVA_AST_IDENTIFIER_EXPR &&
                    strcmp(member->data.member_access.object->data.identifier.name, "Math") == 0) {
                    if (strcmp(method, "abs") == 0 || strcmp(method, "min") == 0 || strcmp(method, "max") == 0) {
                        return rava_type_create_primitive(RAVA_TYPE_INT);
                    } else if (strcmp(method, "round") == 0) {
                        return rava_type_create_primitive(RAVA_TYPE_LONG);
                    } else if (strcmp(method, "sqrt") == 0 || strcmp(method, "pow") == 0 ||
                               strcmp(method, "floor") == 0 || strcmp(method, "ceil") == 0 ||
                               strcmp(method, "sin") == 0 || strcmp(method, "cos") == 0 ||
                               strcmp(method, "tan") == 0 || strcmp(method, "log") == 0 ||
                               strcmp(method, "exp") == 0 || strcmp(method, "random") == 0) {
                        return rava_type_create_primitive(RAVA_TYPE_DOUBLE);
                    }
                }
                if (strcmp(method, "length") == 0) {
                    return rava_type_create_primitive(RAVA_TYPE_INT);
                } else if (strcmp(method, "charAt") == 0) {
                    return rava_type_create_primitive(RAVA_TYPE_CHAR);
                } else if (strcmp(method, "equals") == 0) {
                    return rava_type_create_primitive(RAVA_TYPE_BOOLEAN);
                } else if (strcmp(method, "compareTo") == 0) {
                    return rava_type_create_primitive(RAVA_TYPE_INT);
                } else if (strcmp(method, "size") == 0) {
                    return rava_type_create_primitive(RAVA_TYPE_INT);
                } else if (strcmp(method, "isEmpty") == 0) {
                    return rava_type_create_primitive(RAVA_TYPE_BOOLEAN);
                } else if (strcmp(method, "containsKey") == 0) {
                    return rava_type_create_primitive(RAVA_TYPE_BOOLEAN);
                }
            }
            return rava_type_create_primitive(RAVA_TYPE_INT);
        }

        case RAVA_AST_MEMBER_ACCESS_EXPR: {
            _rava_semantic_check_expression(analyzer, expr->data.member_access.object);
            return rava_type_create_primitive(RAVA_TYPE_INT);
        }

        case RAVA_AST_ARRAY_ACCESS_EXPR: {
            RavaType_t *array_type = _rava_semantic_check_expression(analyzer, expr->data.array_access.array);
            RavaType_t *index_type = _rava_semantic_check_expression(analyzer, expr->data.array_access.index);

            if (index_type && index_type->kind != RAVA_TYPE_INT && index_type->kind != RAVA_TYPE_LONG) {
                _rava_semantic_error(analyzer, "Array index must be an integer", expr->line, expr->column);
            }

            if (array_type && array_type->kind == RAVA_TYPE_ARRAY) {
                return array_type->data.array.element_type;
            }
            return rava_type_create_primitive(RAVA_TYPE_INT);
        }

        case RAVA_AST_NEW_EXPR: {
            RavaType_t *type = _rava_semantic_resolve_type(analyzer, expr->data.new_expr.type);

            for (size_t i = 0; i < expr->data.new_expr.arguments_count; i++) {
                RavaType_t *arg_type = _rava_semantic_check_expression(analyzer, expr->data.new_expr.arguments[i]);
                if (type && type->kind == RAVA_TYPE_ARRAY && arg_type) {
                    if (arg_type->kind != RAVA_TYPE_INT && arg_type->kind != RAVA_TYPE_LONG) {
                        _rava_semantic_error(analyzer, "Array size must be an integer", expr->line, expr->column);
                    }
                }
            }

            return type;
        }

        default:
            return NULL;
    }
}

static bool _rava_semantic_analyze_statement(RavaSemanticAnalyzer_t *analyzer, RavaASTNode_t *stmt) {
    if (!stmt) return true;

    switch (stmt->type) {
        case RAVA_AST_BLOCK_STMT:
            for (size_t i = 0; i < stmt->children_count; i++) {
                _rava_semantic_analyze_node(analyzer, stmt->children[i]);
            }
            return true;

        case RAVA_AST_IF_STMT: {
            RavaType_t *cond_type = _rava_semantic_check_expression(analyzer, stmt->data.if_stmt.condition);
            if (cond_type && cond_type->kind != RAVA_TYPE_BOOLEAN) {
                _rava_semantic_error(analyzer, "If condition must be boolean", stmt->line, stmt->column);
            }
            _rava_semantic_analyze_node(analyzer, stmt->data.if_stmt.then_stmt);
            if (stmt->data.if_stmt.else_stmt) {
                _rava_semantic_analyze_node(analyzer, stmt->data.if_stmt.else_stmt);
            }
            return true;
        }

        case RAVA_AST_WHILE_STMT: {
            RavaType_t *cond_type = _rava_semantic_check_expression(analyzer, stmt->data.while_stmt.condition);
            if (cond_type && cond_type->kind != RAVA_TYPE_BOOLEAN) {
                _rava_semantic_error(analyzer, "While condition must be boolean", stmt->line, stmt->column);
            }
            _rava_semantic_analyze_node(analyzer, stmt->data.while_stmt.body);
            return true;
        }

        case RAVA_AST_FOR_STMT:
            if (stmt->data.for_stmt.init) {
                if (stmt->data.for_stmt.init->type == RAVA_AST_VAR_DECL) {
                    _rava_semantic_analyze_node(analyzer, stmt->data.for_stmt.init);
                } else {
                    _rava_semantic_check_expression(analyzer, stmt->data.for_stmt.init);
                }
            }
            if (stmt->data.for_stmt.condition) {
                RavaType_t *cond_type = _rava_semantic_check_expression(analyzer, stmt->data.for_stmt.condition);
                if (cond_type && cond_type->kind != RAVA_TYPE_BOOLEAN) {
                    _rava_semantic_error(analyzer, "For condition must be boolean", stmt->line, stmt->column);
                }
            }
            if (stmt->data.for_stmt.update) {
                _rava_semantic_check_expression(analyzer, stmt->data.for_stmt.update);
            }
            _rava_semantic_analyze_node(analyzer, stmt->data.for_stmt.body);
            return true;

        case RAVA_AST_RETURN_STMT:
            if (stmt->data.return_stmt.value) {
                _rava_semantic_check_expression(analyzer, stmt->data.return_stmt.value);
            }
            return true;

        case RAVA_AST_EXPR_STMT:
            if (stmt->children_count > 0) {
                _rava_semantic_check_expression(analyzer, stmt->children[0]);
            }
            return true;

        case RAVA_AST_VAR_DECL:
            return _rava_semantic_analyze_var_decl(analyzer, stmt);

        default:
            return true;
    }
}

static bool _rava_semantic_analyze_node(RavaSemanticAnalyzer_t *analyzer, RavaASTNode_t *node) {
    if (!node) return true;

    switch (node->type) {
        case RAVA_AST_COMPILATION_UNIT:
            for (size_t i = 0; i < node->children_count; i++) {
                _rava_semantic_analyze_node(analyzer, node->children[i]);
            }
            return true;

        case RAVA_AST_CLASS_DECL:
            return _rava_semantic_analyze_class(analyzer, node);

        case RAVA_AST_ENUM_DECL:
            return _rava_semantic_analyze_enum(analyzer, node);

        case RAVA_AST_METHOD_DECL:
            return _rava_semantic_analyze_method(analyzer, node);

        case RAVA_AST_FIELD_DECL:
            return _rava_semantic_analyze_field(analyzer, node);

        case RAVA_AST_CONSTRUCTOR_DECL:
            return true;

        default:
            return _rava_semantic_analyze_statement(analyzer, node);
    }
}

bool rava_semantic_analyze(RavaSemanticAnalyzer_t *analyzer, RavaASTNode_t *root) {
    if (!analyzer || !root) return false;

    _rava_semantic_analyze_node(analyzer, root);

    return !analyzer->had_error;
}

RavaType_t* rava_semantic_get_expression_type(RavaSemanticAnalyzer_t *analyzer, RavaASTNode_t *expr) {
    return _rava_semantic_check_expression(analyzer, expr);
}