#define _POSIX_C_SOURCE 200809L
#include "repl_session.h"
#include <stdlib.h>
#include <string.h>
#include <stdio.h>

RavaREPLSession_t* rava_repl_session_create(void) {
    RavaREPLSession_t *session = calloc(1, sizeof(RavaREPLSession_t));
    if (!session) return NULL;

    session->variable_capacity = RAVA_REPL_MAX_VARIABLES;
    session->variables = calloc(session->variable_capacity, sizeof(RavaREPLVariable_t));
    if (!session->variables) {
        free(session);
        return NULL;
    }

    session->method_capacity = RAVA_REPL_MAX_METHODS;
    session->methods = calloc(session->method_capacity, sizeof(RavaREPLMethod_t));
    if (!session->methods) {
        free(session->variables);
        free(session);
        return NULL;
    }

    session->class_capacity = RAVA_REPL_MAX_CLASSES;
    session->classes = calloc(session->class_capacity, sizeof(RavaREPLClass_t));
    if (!session->classes) {
        free(session->methods);
        free(session->variables);
        free(session);
        return NULL;
    }

    session->history = rava_repl_history_create(RAVA_REPL_MAX_HISTORY);
    session->start_time = time(NULL);
    session->debug_mode = false;
    session->execution_count = 0;
    session->program = NULL;
    session->session_class = NULL;
    session->vm = NULL;
    session->last_error = NULL;

    return session;
}

void rava_repl_session_destroy(RavaREPLSession_t *session) {
    if (!session) return;

    for (size_t i = 0; i < session->variable_count; i++) {
        free(session->variables[i].name);
        free(session->variables[i].type_name);
        free(session->variables[i].initializer);
    }
    free(session->variables);

    for (size_t i = 0; i < session->method_count; i++) {
        free(session->methods[i].name);
        free(session->methods[i].return_type);
        free(session->methods[i].source);
        for (size_t j = 0; j < session->methods[i].param_count; j++) {
            free(session->methods[i].param_names[j]);
            free(session->methods[i].param_types[j]);
        }
        free(session->methods[i].param_names);
        free(session->methods[i].param_types);
    }
    free(session->methods);

    for (size_t i = 0; i < session->class_count; i++) {
        free(session->classes[i].name);
        free(session->classes[i].source);
        for (size_t j = 0; j < session->classes[i].field_count; j++) {
            free(session->classes[i].field_names[j]);
            free(session->classes[i].field_types[j]);
        }
        free(session->classes[i].field_names);
        free(session->classes[i].field_types);
        for (size_t j = 0; j < session->classes[i].method_count; j++) {
            free(session->classes[i].method_names[j]);
        }
        free(session->classes[i].method_names);
    }
    free(session->classes);

    rava_repl_history_destroy(session->history);

    if (session->vm) {
        rava_vm_destroy(session->vm);
    }
    if (session->program) {
        rava_program_destroy(session->program);
    }
    free(session->last_error);
    free(session);
}

bool rava_repl_session_add_variable(RavaREPLSession_t *session, const char *name,
                                     const char *type_name, const char *initializer,
                                     RavaValueType_e type, RavaValue_t value) {
    if (!session || !name) return false;

    RavaREPLVariable_t *existing = rava_repl_session_get_variable(session, name);
    if (existing) {
        free(existing->type_name);
        free(existing->initializer);
        existing->type_name = type_name ? strdup(type_name) : NULL;
        existing->initializer = initializer ? strdup(initializer) : NULL;
        existing->type = type;
        existing->value = value;
        existing->is_initialized = true;
        return true;
    }

    if (session->variable_count >= session->variable_capacity) {
        return false;
    }

    RavaREPLVariable_t *var = &session->variables[session->variable_count++];
    var->name = strdup(name);
    var->type_name = type_name ? strdup(type_name) : NULL;
    var->initializer = initializer ? strdup(initializer) : NULL;
    var->type = type;
    var->value = value;
    var->is_initialized = true;

    return true;
}

RavaREPLVariable_t* rava_repl_session_get_variable(RavaREPLSession_t *session, const char *name) {
    if (!session || !name) return NULL;
    for (size_t i = 0; i < session->variable_count; i++) {
        if (strcmp(session->variables[i].name, name) == 0) {
            return &session->variables[i];
        }
    }
    return NULL;
}

bool rava_repl_session_set_variable(RavaREPLSession_t *session, const char *name, RavaValue_t value) {
    RavaREPLVariable_t *var = rava_repl_session_get_variable(session, name);
    if (!var) return false;
    var->value = value;
    var->type = value.type;
    return true;
}

void rava_repl_session_clear_variables(RavaREPLSession_t *session) {
    if (!session) return;
    for (size_t i = 0; i < session->variable_count; i++) {
        free(session->variables[i].name);
        free(session->variables[i].type_name);
        free(session->variables[i].initializer);
    }
    session->variable_count = 0;
}

bool rava_repl_session_add_method(RavaREPLSession_t *session, const char *name,
                                   const char *return_type, const char *source) {
    if (!session || !name) return false;

    RavaREPLMethod_t *existing = rava_repl_session_get_method(session, name);
    if (existing) {
        free(existing->return_type);
        free(existing->source);
        existing->return_type = return_type ? strdup(return_type) : NULL;
        existing->source = source ? strdup(source) : NULL;
        existing->compiled = NULL;
        return true;
    }

    if (session->method_count >= session->method_capacity) {
        return false;
    }

    RavaREPLMethod_t *method = &session->methods[session->method_count++];
    method->name = strdup(name);
    method->return_type = return_type ? strdup(return_type) : NULL;
    method->source = source ? strdup(source) : NULL;
    method->param_names = NULL;
    method->param_types = NULL;
    method->param_count = 0;
    method->compiled = NULL;

    return true;
}

RavaREPLMethod_t* rava_repl_session_get_method(RavaREPLSession_t *session, const char *name) {
    if (!session || !name) return NULL;
    for (size_t i = 0; i < session->method_count; i++) {
        if (strcmp(session->methods[i].name, name) == 0) {
            return &session->methods[i];
        }
    }
    return NULL;
}

bool rava_repl_session_add_class(RavaREPLSession_t *session, const char *name, const char *source) {
    if (!session || !name) return false;

    RavaREPLClass_t *existing = rava_repl_session_get_class(session, name);
    if (existing) {
        free(existing->source);
        existing->source = source ? strdup(source) : NULL;
        existing->compiled = NULL;
        return true;
    }

    if (session->class_count >= session->class_capacity) {
        return false;
    }

    RavaREPLClass_t *cls = &session->classes[session->class_count++];
    cls->name = strdup(name);
    cls->source = source ? strdup(source) : NULL;
    cls->field_names = NULL;
    cls->field_types = NULL;
    cls->field_count = 0;
    cls->method_names = NULL;
    cls->method_count = 0;
    cls->compiled = NULL;

    return true;
}

RavaREPLClass_t* rava_repl_session_get_class(RavaREPLSession_t *session, const char *name) {
    if (!session || !name) return NULL;
    for (size_t i = 0; i < session->class_count; i++) {
        if (strcmp(session->classes[i].name, name) == 0) {
            return &session->classes[i];
        }
    }
    return NULL;
}

char* rava_repl_session_generate_source(RavaREPLSession_t *session) {
    if (!session) return NULL;

    size_t capacity = 4096;
    char *source = malloc(capacity);
    if (!source) return NULL;
    size_t len = 0;
    source[0] = '\0';

#define APPEND(...) do { \
    int needed = snprintf(source + len, capacity - len, __VA_ARGS__); \
    if (needed < 0) { free(source); return NULL; } \
    if ((size_t)needed >= capacity - len) { \
        capacity = capacity * 2 + needed; \
        char *tmp = realloc(source, capacity); \
        if (!tmp) { free(source); return NULL; } \
        source = tmp; \
        snprintf(source + len, capacity - len, __VA_ARGS__); \
    } \
    len += needed; \
} while(0)

    for (size_t i = 0; i < session->class_count; i++) {
        if (session->classes[i].source) {
            APPEND("%s\n\n", session->classes[i].source);
        }
    }

    APPEND("public class _REPL_Session_ {\n");

    for (size_t i = 0; i < session->variable_count; i++) {
        RavaREPLVariable_t *var = &session->variables[i];
        const char *type = var->type_name ? var->type_name : "int";
        APPEND("    public static %s %s;\n", type, var->name);
    }

    if (session->variable_count > 0) {
        APPEND("\n");
    }

    for (size_t i = 0; i < session->method_count; i++) {
        if (session->methods[i].source) {
            APPEND("    %s\n\n", session->methods[i].source);
        }
    }

    APPEND("}\n");

#undef APPEND

    return source;
}

bool rava_repl_session_recompile(RavaREPLSession_t *session) {
    if (!session) return false;

    char *source = rava_repl_session_generate_source(session);
    if (!source) {
        rava_repl_session_set_error(session, "Failed to generate session source");
        return false;
    }

    if (session->debug_mode) {
        fprintf(stderr, "[DEBUG] Session source:\n%s\n", source);
    }

    RavaLexer_t *lexer = rava_lexer_create(source);
    RavaParser_t *parser = rava_parser_create(lexer);
    RavaASTNode_t *ast = rava_parser_parse(parser);

    if (!ast || parser->had_error) {
        rava_repl_session_set_error(session, parser->error_message ? parser->error_message : "Parse error");
        rava_parser_destroy(parser);
        rava_lexer_destroy(lexer);
        free(source);
        return false;
    }

    RavaSemanticAnalyzer_t *analyzer = rava_semantic_analyzer_create();
    if (!rava_semantic_analyze(analyzer, ast)) {
        rava_repl_session_set_error(session, analyzer->error_message ? analyzer->error_message : "Semantic error");
        rava_ast_node_destroy(ast);
        rava_parser_destroy(parser);
        rava_lexer_destroy(lexer);
        rava_semantic_analyzer_destroy(analyzer);
        free(source);
        return false;
    }

    RavaIRGenerator_t *ir_gen = rava_ir_generator_create(analyzer);
    RavaProgram_t *program = rava_ir_generate(ir_gen, ast);

    if (!program) {
        rava_repl_session_set_error(session, "IR generation failed");
        rava_ast_node_destroy(ast);
        rava_parser_destroy(parser);
        rava_lexer_destroy(lexer);
        rava_semantic_analyzer_destroy(analyzer);
        rava_ir_generator_destroy(ir_gen);
        free(source);
        return false;
    }

    if (session->vm) {
        rava_vm_destroy(session->vm);
    }
    if (session->program) {
        rava_program_destroy(session->program);
    }

    session->program = program;
    session->vm = rava_vm_create(program);
    session->session_class = NULL;

    rava_ast_node_destroy(ast);
    rava_parser_destroy(parser);
    rava_lexer_destroy(lexer);
    rava_semantic_analyzer_destroy(analyzer);
    rava_ir_generator_destroy(ir_gen);
    free(source);

    return true;
}

void rava_repl_session_set_error(RavaREPLSession_t *session, const char *error) {
    if (!session) return;
    free(session->last_error);
    session->last_error = error ? strdup(error) : NULL;
}

const char* rava_repl_session_get_error(RavaREPLSession_t *session) {
    if (!session) return NULL;
    return session->last_error;
}