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

#include "spl_comp.h"
#include <stdarg.h>
#include <stdio.h>
#include <stdlib.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);
    vec_init(ctx->ns_chain);
    map_init(ctx->type_defs, MAP_HASH_STR, MAP_CMP_STR);
    map_init(ctx->const_values, MAP_HASH_STR, MAP_CMP_STR);
    vec_init(ctx->call_fixups);
    vec_init(ctx->call_fixup_funcs);
    spl_prog_init(&ctx->prog);
    ctx->error_msg[0] = '\0';
}

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_free(ctx->funcs);
    vec_free(ctx->ns_chain);
    /* Free type defs — complex, leak for now in bootstrap */
    map_free(ctx->type_defs);
    map_free(ctx->const_values);
    vec_free(ctx->call_fixups);
    vec_free(ctx->call_fixup_funcs);
    spl_prog_drop(&ctx->prog);
    free(ctx->break_patches);
}

void spl_comp_reset(spl_comp_t *ctx) {
    /* Keep prog, reset everything else */
    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);
    vec_free(ctx->ns_chain);
    vec_init(ctx->ns_chain);
    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 = NULL;
    ctx->current_local_bytes = 0;
    ctx->peak_local_bytes = 0;
    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->call_fixups);
    vec_init(ctx->call_fixups);
    vec_free(ctx->call_fixup_funcs);
    vec_init(ctx->call_fixup_funcs);
}

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;
}

/* ---- Codegen helpers ---- */

spl_val_t spl_emit(spl_comp_t *ctx, uint16_t opcode, uint16_t type, spl_val_t imm) {
    return spl_prog_emit(&ctx->prog, opcode, type, imm);
}

void spl_patch(spl_comp_t *ctx, spl_val_t addr, spl_val_t target) {
    if (addr < vec_size(ctx->prog.insns)) {
        vec_at(ctx->prog.insns, addr).imm = target;
    }
}

spl_val_t spl_emit_jmp(spl_comp_t *ctx) {
    /* Emit JMP with placeholder 0, return address to patch */
    return spl_emit(ctx, SPL_JMP, SPL_VOID, 0);
}

spl_val_t spl_emit_bz(spl_comp_t *ctx) { return spl_emit(ctx, SPL_BZ, SPL_VOID, 0); }

spl_val_t spl_emit_bnz(spl_comp_t *ctx) { return spl_emit(ctx, SPL_BNZ, SPL_VOID, 0); }

void spl_patch_to_here(spl_comp_t *ctx, spl_val_t addr) {
    /* Compute relative offset: target - (source + 1) */
    spl_val_t here = vec_size(ctx->prog.insns);
    spl_val_t offset = here - addr - 1;
    spl_patch(ctx, addr, offset);
}

/* ---- Multi-slot copy helpers ---- */

void spl_emit_copy_slots(spl_comp_t *ctx, int dest_offset, usize nslots) {
    for (usize i = 0; i < nslots; i++) {
        if (i < nslots - 1)
            spl_emit(ctx, SPL_DUP, SPL_VOID, 0);
        if (i > 0) {
            spl_emit(ctx, SPL_PUSH, SPL_USIZE, i * sizeof(spl_val_t));
            spl_emit(ctx, SPL_ADD, SPL_USIZE, 0);
        }
        spl_emit(ctx, SPL_LOAD, SPL_PTR, 0);
        spl_emit(ctx, SPL_LADDR, SPL_PTR, dest_offset + (int)(i * sizeof(spl_val_t)));
        spl_emit(ctx, SPL_SWAP, SPL_VOID, 0);
        spl_emit(ctx, SPL_STORE, SPL_PTR, 0);
    }
}

/* Copy nslots from source address to destination address.
 * Stack before: [..., depth_items..., dest_addr, src_addr] (src TOS)
 * Stack after: [..., depth_items...]
 * depth: number of items below dest_addr that must be preserved (e.g. 1 if
 *        there's a base pointer between the rest of the stack and dest_addr). */
void spl_emit_copy_addr_to_addr(spl_comp_t *ctx, usize nslots, int depth) {
    for (usize i = 0; i < nslots; i++) {
        spl_emit(ctx, SPL_DUP, SPL_VOID, 0);
        if (i > 0) {
            spl_emit(ctx, SPL_PUSH, SPL_USIZE, i * sizeof(spl_val_t));
            spl_emit(ctx, SPL_ADD, SPL_USIZE, 0);
        }
        spl_emit(ctx, SPL_LOAD, SPL_PTR, 0);
        spl_emit(ctx, SPL_PICK, SPL_VOID, 2 + depth);
        if (i > 0) {
            spl_emit(ctx, SPL_PUSH, SPL_USIZE, i * sizeof(spl_val_t));
            spl_emit(ctx, SPL_ADD, SPL_USIZE, 0);
        }
        spl_emit(ctx, SPL_SWAP, SPL_VOID, 0);
        spl_emit(ctx, SPL_STORE, SPL_PTR, 0);
    }
    spl_emit(ctx, SPL_DROP, SPL_VOID, 0);
    spl_emit(ctx, SPL_DROP, SPL_VOID, 0);
}

/* ---- Uniform LOAD/STORE type helper ---- */

spl_type_t spl_type_emit_type(spl_type_info_t *type) {
    if (!type)
        return SPL_I32;
    if (type->kind == TYPE_BASIC)
        return type->basic_type;
    if (type->kind == TYPE_PTR)
        return SPL_PTR;
    if (type->kind == TYPE_ENUM) {
        /* Simple enum (no data variants): tag is I32 */
        vec_for(type->variants, i) {
            if (vec_at(type->variants, i).data_type)
                return SPL_PTR; /* has data — treat as aggregate */
        }
        return SPL_I32;
    }
    return SPL_PTR;
}

/* Scalar types fit in one slot, are auto-loaded on field access,
 * and support == / != comparison directly. */
int spl_type_is_scalar(spl_type_info_t *type) {
    if (!type)
        return 1;
    if (type->kind == TYPE_BASIC || type->kind == TYPE_PTR)
        return 1;
    if (type->kind == TYPE_ENUM) {
        vec_for(type->variants, i) {
            if (vec_at(type->variants, i).data_type)
                return 0; /* has data — aggregate */
        }
        return 1; /* simple enum — scalar */
    }
    return 0;
}

/* ---- Load from [saved_ptr+offset], store to local var ---- */

void spl_emit_load_to_var(spl_comp_t *ctx, int ptr_slot_offset, usize byte_offset,
                          spl_type_info_t *data_type, int var_offset) {
    spl_emit(ctx, SPL_LADDR, SPL_PTR, ptr_slot_offset);
    spl_emit(ctx, SPL_LOAD, SPL_PTR, 0);
    if (byte_offset > 0) {
        spl_emit(ctx, SPL_PUSH, SPL_USIZE, byte_offset);
        spl_emit(ctx, SPL_ADD, SPL_USIZE, 0);
    }
    if (data_type && !spl_type_is_scalar(data_type) &&
        spl_type_size(data_type) > sizeof(spl_val_t)) {
        spl_emit_copy_slots(ctx, var_offset, spl_type_slot_count(data_type));
    } else {
        spl_type_t bt = spl_type_emit_type(data_type);
        spl_emit(ctx, SPL_LOAD, bt, 0);
        spl_emit(ctx, SPL_LADDR, SPL_PTR, var_offset);
        spl_emit(ctx, SPL_SWAP, SPL_VOID, 0);
        spl_emit(ctx, SPL_STORE, bt, 0);
    }
}

/* ---- 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--;
    }
}

/* ---- Namespace chain management ---- */

void spl_ns_push(spl_comp_t *ctx, spl_type_info_t *type) {
    if (type)
        vec_push(ctx->ns_chain, type);
}

void spl_ns_pop(spl_comp_t *ctx) {
    if (vec_size(ctx->ns_chain) > 0)
        ctx->ns_chain.size--;
}

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

int spl_declare_var(spl_comp_t *ctx, const char *name, spl_type_info_t *type, int is_const) {
    spl_var_info_t var;
    memset(&var, 0, sizeof(var));
    var.name = strdup(name);
    var.type = type;
    var.is_const = is_const;
    var.depth = ctx->scope_depth;
    var.offset = ctx->current_local_bytes;

    usize var_size = spl_type_size(type);
    /* Round up to sizeof(spl_val_t) alignment so params (pushed as spl_val_t) align */
    usize aligned = (var_size + sizeof(spl_val_t) - 1) & ~(sizeof(spl_val_t) - 1);
    if (aligned < sizeof(spl_val_t))
        aligned = sizeof(spl_val_t);
    ctx->current_local_bytes += (int)aligned;
    if (ctx->current_local_bytes > ctx->peak_local_bytes)
        ctx->peak_local_bytes = ctx->current_local_bytes;

    /* Add to current scope */
    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);
        /* Search in reverse order so later declarations shadow earlier ones */
        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, spl_type_info_t *ret_type, int nparams,
                     int is_extern, int is_pub) {
    /* Check if already declared — allows pre-registration */
    vec_for(ctx->funcs, i) {
        spl_func_info_t *existing = &vec_at(ctx->funcs, i);
        if (strcmp(existing->name, name) == 0) {
            existing->ret_type = ret_type;
            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 = ret_type;
    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) {
    /* 1. Walk ns_chain from innermost to outermost, check each type's methods */
    for (usize i = vec_size(ctx->ns_chain); i > 0; i--) {
        spl_type_info_t *type = vec_at(ctx->ns_chain, i - 1);
        vec_for(type->methods, j) {
            if (strcmp(vec_at(type->methods, j).name, name) == 0)
                return vec_at(type->methods, j).func_idx;
        }
    }

    /* 2. Fallback to flat func table (top-level functions and qualified names) */
    vec_for(ctx->funcs, j) {
        if (strcmp(vec_at(ctx->funcs, j).name, name) == 0)
            return (int)j;
    }
    return -1;
}

/* Ensure a native function is registered for NCALL dispatch.
 * Returns the native index. */
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;

    /* Emit JMP placeholder (will skip defer body during normal execution).
     * This JMP is at the current position. The defer body starts right after. */
    spl_val_t jmp_skip = spl_emit_jmp(ctx);

    spl_defer_entry_t *e = &ctx->defer_stack[ctx->defer_count];
    e->body_start = jmp_skip + 1; /* instruction right after the JMP = body start */
    e->jmp_exit = 0;              /* set after body is parsed */
    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) {
    /* === Pass 1: patch skip JMPs to jump past their defer body ===
     * During normal execution, the skip JMP at body_start-1 must jump
     * over the defer body (to jmp_exit + 1 = the code after the defer). */
    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; /* the JMP L_skip instruction */
        spl_val_t skip_target = e->jmp_exit + 1; /* instruction right after defer body */
        spl_val_t skip_offset = skip_target - skip_addr - 1;
        spl_patch(ctx, skip_addr, skip_offset);
    }

    /* === Pass 2: at scope exit, emit backwards JMPs to each defer body ===
     * Process in reverse order so the LAST declared defer runs FIRST at scope exit.
     * Each defer body's trailing JMP_exit gets patched to jump to right after
     * the scope-exit JMP we just emitted, so control chains through correctly. */
    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];

        /* Emit JMP backwards to the defer body */
        spl_val_t here = vec_size(ctx->prog.insns);
        spl_val_t jmp_offset = (spl_val_t)((isize)e->body_start - (isize)here - 1);
        spl_emit(ctx, SPL_JMP, SPL_VOID, jmp_offset);

        /* Patch the body's trailing JMP_exit to jump to here+1 (right after the
         * backwards JMP we just emitted). This chains to the next defer or exits. */
        if (e->jmp_exit > 0) {
            spl_val_t exit_target = here + 1;
            spl_val_t offset = exit_target - e->jmp_exit - 1;
            spl_patch(ctx, e->jmp_exit, offset);
        }
    }

    /* Remove processed defers from stack */
    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) {
    /* Runtime support functions for compiled SPL programs go here.
     * For stage1 bootstrap, most operations are handled inline.
     * We register basic runtime helpers if needed. */
    (void)prog;
}

/* ---- Patch call fixups ---- */

void spl_patch_call_fixups(spl_comp_t *ctx) {
    for (usize i = 0; i < vec_size(ctx->call_fixups); i++) {
        spl_val_t insn_idx = vec_at(ctx->call_fixups, i);
        int pfi = vec_at(ctx->call_fixup_funcs, i);
        if (pfi >= 0 && pfi < (int)vec_size(ctx->prog.funcs)) {
            spl_val_t addr = vec_at(ctx->prog.funcs, pfi).address;
            vec_at(ctx->prog.insns, insn_idx).imm = addr;
        }
    }
}

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

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

    /* Skip initial newlines */
    while (ctx->tok_idx < vec_size(ctx->toks) &&
           vec_at(ctx->toks, ctx->tok_idx).type == TOK_ENDLINE)
        ctx->tok_idx++;

    /* Phase 2-4: Parse and codegen */
    spl_parse_prog(ctx);
    if (ctx->has_error)
        return -1;

    /* Phase 5: Patch all function call fixups (forward references) */
    spl_patch_call_fixups(ctx);

    return 0;
}