/* spl_comp.c — SPL compiler main logic and codegen helpers */

#include "spl_comp.h"
#include "spl_lex_util.h"
#include <stdarg.h>
#include <stdio.h>
#include <string.h>

/* ---- Compiler context init/drop ---- */

void spl_comp_init(spl_comp_t *ctx) {
    memset(ctx, 0, sizeof(*ctx));
    vec_init(ctx->toks);
    vec_init(ctx->scopes);
    vec_init(ctx->funcs);
    map_init(ctx->const_values, MAP_HASH_STR, MAP_CMP_STR);
    spl_prog_init(&ctx->prog);
    ctx->error_msg[0] = '\0';
    spl_emit_init(&ctx->emit, &ctx->prog);
    spl_type_ctx_init(&ctx->tctx);
    ctx->current_ret_type_idx = -1;
}

void spl_comp_drop(spl_comp_t *ctx) {
    if (!ctx)
        return;
    spl_tok_vec_drop(&ctx->toks);
    vec_for(ctx->scopes, i) { vec_free(vec_at(ctx->scopes, i).vars); }
    vec_free(ctx->scopes);
    vec_for(ctx->funcs, i) {
        free(ctx->funcs.data[i].name);
        free(ctx->funcs.data[i].param_type_indices);
        free(ctx->funcs.data[i].param_names);
    }
    vec_free(ctx->funcs);
    map_free(ctx->const_values);
    spl_prog_drop(&ctx->prog);
    free(ctx->break_patches);
    spl_emit_drop(&ctx->emit);
    spl_type_ctx_drop(&ctx->tctx);
}

void spl_comp_reset(spl_comp_t *ctx) {
    spl_tok_vec_drop(&ctx->toks);
    vec_init(ctx->toks);
    vec_for(ctx->scopes, i) { vec_free(vec_at(ctx->scopes, i).vars); }
    vec_free(ctx->scopes);
    vec_init(ctx->scopes);
    ctx->scope_depth = 0;
    ctx->tok_idx = 0;
    ctx->has_error = 0;
    ctx->error_msg[0] = '\0';
    ctx->current_func_idx = -1;
    ctx->current_ret_type_idx = -1;
    fa_init(&ctx->emit.frame);
    ctx->in_loop = 0;
    ctx->break_patch_count = 0;
    ctx->break_patch_cap = 0;
    ctx->continue_target = 0;
    ctx->defer_count = 0;
    ctx->next_gdata_idx = 0;
    ctx->addr_of_mode = 0;
    free(ctx->break_patches);
    ctx->break_patches = NULL;
    vec_free(ctx->emit.fixups);
    vec_init(ctx->emit.fixups);
}

void spl_comp_error(spl_comp_t *ctx, const char *fmt, ...) {
    va_list args;
    va_start(args, fmt);
    vsnprintf(ctx->error_msg, COMP_ERROR_MAX - 1, fmt, args);
    va_end(args);
    ctx->has_error = 1;
}

/* ---- Scope management ---- */

void spl_push_scope(spl_comp_t *ctx) {
    spl_scope_t scope;
    vec_init(scope.vars);
    scope.depth = ++ctx->scope_depth;
    vec_push(ctx->scopes, scope);
}

void spl_pop_scope(spl_comp_t *ctx) {
    if (vec_size(ctx->scopes) > 0) {
        spl_scope_t *scope = &vec_at(ctx->scopes, vec_size(ctx->scopes) - 1);
        vec_for(scope->vars, i) { free(vec_at(scope->vars, i).name); }
        vec_free(scope->vars);
        ctx->scope_depth--;
        ctx->scopes.size--;
    } else {
        ctx->scope_depth--;
    }
}

/* ---- Variable management ---- */

int spl_declare_var(spl_comp_t *ctx, const char *name, int type_idx, int is_const) {
    spl_var_info_t var;
    memset(&var, 0, sizeof(var));
    var.name = strdup(name);
    var.type_idx = type_idx;
    var.is_const = is_const;
    var.depth = ctx->scope_depth;
    var.offset = fa_alloc(&ctx->emit.frame, spl_type_size(&ctx->tctx, type_idx));

    if (vec_size(ctx->scopes) > 0) {
        spl_scope_t *scope = &vec_at(ctx->scopes, vec_size(ctx->scopes) - 1);
        vec_push(scope->vars, var);
    }
    return var.offset;
}

spl_var_info_t *spl_lookup_var(spl_comp_t *ctx, const char *name) {
    for (int i = (int)vec_size(ctx->scopes) - 1; i >= 0; i--) {
        spl_scope_t *scope = &vec_at(ctx->scopes, i);
        for (int j = (int)vec_size(scope->vars) - 1; j >= 0; j--) {
            if (strcmp(vec_at(scope->vars, j).name, name) == 0)
                return &vec_at(scope->vars, j);
        }
    }
    return NULL;
}

int spl_get_var_offset(spl_comp_t *ctx, const char *name) {
    spl_var_info_t *v = spl_lookup_var(ctx, name);
    return v ? v->offset : -1;
}

/* ---- Function management ---- */

int spl_declare_func(spl_comp_t *ctx, const char *name, int ret_type_idx, int nparams,
                     int is_extern, int is_pub) {
    vec_for(ctx->funcs, i) {
        spl_func_info_t *existing = &vec_at(ctx->funcs, i);
        if (strcmp(existing->name, name) == 0) {
            existing->ret_type_idx = ret_type_idx;
            existing->nparams = nparams;
            existing->is_extern = is_extern;
            existing->is_pub = is_pub;
            existing->func_idx =
                spl_prog_update_func(&ctx->prog, existing->func_idx, name, nparams);
            return (int)i;
        }
    }

    spl_func_info_t fi;
    memset(&fi, 0, sizeof(fi));
    fi.name = strdup(name);
    fi.ret_type_idx = ret_type_idx;
    fi.nparams = nparams;
    fi.is_extern = is_extern;
    fi.is_pub = is_pub;

    if (is_extern)
        fi.func_idx = -1;
    else
        fi.func_idx = spl_prog_add_func_simple(&ctx->prog, name, nparams);

    vec_push(ctx->funcs, fi);
    return (int)vec_size(ctx->funcs) - 1;
}

int spl_lookup_func(spl_comp_t *ctx, const char *name) {
    {
        int current = ctx->tctx.current_type_idx;
        while (current >= 0) {
            spl_type_item_vec_t *items = spl_type_items(&ctx->tctx, current);
            if (items) {
                vec_for(*items, j) {
                    spl_type_item_t *it = &vec_at(*items, j);
                    if (it->item_kind == ITEM_METHOD && it->name && strcmp(it->name, name) == 0)
                        return it->method.func_idx;
                }
            }
            current = vec_at(ctx->tctx.types, current).parent_type_idx;
        }
    }

    vec_for(ctx->funcs, j) {
        if (strcmp(vec_at(ctx->funcs, j).name, name) == 0)
            return (int)j;
    }
    return -1;
}

int spl_ensure_native(spl_comp_t *ctx, const char *name) {
    vec_for(ctx->prog.natives, ni) {
        if (strcmp(vec_at(ctx->prog.natives, ni).name, name) == 0)
            return (int)ni;
    }
    spl_native_t nat;
    nat.name = strdup(name);
    nat.idx_of_strtab = 0;
    nat.impl_fn = NULL;
    vec_push(ctx->prog.natives, nat);
    return (int)vec_size(ctx->prog.natives) - 1;
}

/* ---- String/data management ---- */

int spl_add_string(spl_comp_t *ctx, const char *str) {
    return spl_add_global_data(ctx, (void *)str, strlen(str) + 1);
}

int spl_add_global_data(spl_comp_t *ctx, void *data, usize size) {
    return spl_prog_add_data(&ctx->prog, data, size);
}

/* ---- Defer ---- */

void spl_emit_defer(spl_comp_t *ctx) {
    if (ctx->defer_count >= DEFER_MAX)
        return;

    spl_val_t jmp_skip = emit_jmp_here(&ctx->emit);

    spl_defer_entry_t *e = &ctx->defer_stack[ctx->defer_count];
    e->body_start = jmp_skip + 1;
    e->jmp_exit = 0;
    e->depth = ctx->scope_depth;
    e->count_at_decl = ctx->defer_count;
    ctx->defer_count++;
}

void spl_emit_defer_epilogue(spl_comp_t *ctx, int depth) {
    for (int i = 0; i < ctx->defer_count; i++) {
        if (ctx->defer_stack[i].depth != depth)
            continue;
        spl_defer_entry_t *e = &ctx->defer_stack[i];

        spl_val_t skip_addr = e->body_start - 1;
        spl_val_t skip_target = e->jmp_exit + 1;
        emit_patch(&ctx->emit, skip_addr, skip_target - skip_addr - 1);
    }

    for (int i = ctx->defer_count - 1; i >= 0; i--) {
        if (ctx->defer_stack[i].depth != depth)
            continue;
        spl_defer_entry_t *e = &ctx->defer_stack[i];

        spl_val_t here = vec_size(ctx->prog.insns);
        spl_val_t jmp_offset = (spl_val_t)((isize)e->body_start - (isize)here - 1);
        emit_jmp(&ctx->emit, jmp_offset);

        if (e->jmp_exit > 0)
            emit_patch(&ctx->emit, e->jmp_exit, here + 1 - e->jmp_exit - 1);
    }

    int new_count = 0;
    for (int i = 0; i < ctx->defer_count; i++) {
        if (ctx->defer_stack[i].depth != depth)
            ctx->defer_stack[new_count++] = ctx->defer_stack[i];
    }
    ctx->defer_count = new_count;
}

/* ---- Register runtime natives ---- */

void spl_comp_register(spl_prog_t *prog) { (void)prog; }

/* ---- Main compilation ---- */

int spl_compile(spl_comp_t *ctx, const char *source, const char *fname) {
    ctx->toks = spl_lex(source, fname);
    ctx->tok_idx = 0;
    ctx->fname = fname;
    ctx->source = source;

    while (ctx->tok_idx < vec_size(ctx->toks) &&
           vec_at(ctx->toks, ctx->tok_idx).type == TOK_ENDLINE)
        ctx->tok_idx++;

    spl_parse_prog(ctx);
    if (ctx->has_error)
        return -1;

    emit_patch_call_fixups(&ctx->emit, &ctx->prog);

    return 0;
}