/* spl_stmt.c — Statement parser + codegen */

#include "spl_comp.h"
#include "spl_lex_util.h"
#include <stdlib.h>
#include <string.h>

/* ============================================================
 * Return statement: ret expr;
 * ============================================================ */

static void parse_ret_stmt(spl_comp_t *ctx) {
    advance(ctx); /* ret */
    skip_nl(ctx);

    /* Before returning, execute all pending defers from innermost scope outward */
    for (int d = ctx->scope_depth; d >= 1; d--) {
        spl_emit_defer_epilogue(ctx, d);
    }

    if (peek(ctx)->type == TOK_SEMICOLON || peek(ctx)->type == TOK_R_BRACE ||
        peek(ctx)->type == TOK_ENDLINE) {
        /* void return */
        spl_emit(ctx, SPL_RET, SPL_VOID, 0);
    } else {
        spl_expr_result_t val = spl_parse_expr(ctx, PREC_MIN);
        (void)val;
        spl_type_t rt = SPL_VOID;
        if (ctx->current_ret_type) {
            if (ctx->current_ret_type->kind == TYPE_BASIC) {
                rt = ctx->current_ret_type->basic_type;
            } else if (ctx->current_ret_type->kind == TYPE_PTR) {
                rt = SPL_PTR;
            } else if (spl_type_size(ctx->current_ret_type) <= sizeof(spl_val_t)) {
                /* 1-slot struct/enum/etc: use PTR type */
                rt = SPL_PTR;
            }
        }
        spl_emit(ctx, SPL_RET, rt, 0);
    }
    if (peek(ctx)->type == TOK_SEMICOLON)
        advance(ctx);
}

/* ============================================================
 * Variable declaration: var name: Type [= expr];
 * ============================================================ */

static void parse_var_decl(spl_comp_t *ctx, int is_const) {
    advance(ctx); /* var or const */
    skip_nl(ctx);

    spl_tok_t *name_tok = advance(ctx);
    char vname[256];
    usize nlen = name_tok->len < 255 ? name_tok->len : 255;
    memcpy(vname, name_tok->lexeme, nlen);
    vname[nlen] = '\0';

    spl_type_info_t *var_type = NULL;
    int has_init = 0;

    skip_nl(ctx);

    if (peek(ctx)->type == TOK_COLON) {
        advance(ctx); /* : */
        skip_nl(ctx);
        /* Check for := */
        if (peek(ctx)->type == TOK_ASSIGN) {
            /* := is colon-assign */
            has_init = 1;
        } else {
            var_type = spl_parse_type(ctx);
            skip_nl(ctx);
        }
    }

    if (peek(ctx)->type == TOK_COLON_ASSIGN || peek(ctx)->type == TOK_ASSIGN) {
        has_init = 1;
        if (peek(ctx)->type == TOK_COLON_ASSIGN)
            advance(ctx);
        else
            advance(ctx); /* = */
    }

    /* Init expression: parse before declare to enable type inference */
    spl_expr_result_t init = {0};
    if (has_init) {
        skip_nl(ctx);
        init = spl_parse_expr(ctx, PREC_MIN);
    }

    if (!var_type) {
        var_type = init.type ? init.type : spl_type_basic(SPL_I32);
    }
    int offset = spl_declare_var(ctx, vname, var_type, is_const);

    /* Store init value */
    if (has_init) {
        /* Handle slice struct literal: { .ptr = ..., .len = ... } */
        if (var_type && var_type->kind == TYPE_SLICE && peek(ctx)->type == TOK_L_BRACE && !init.type) {
            advance(ctx); /* { */
            skip_nl(ctx);
            while (peek(ctx)->type != TOK_R_BRACE && peek(ctx)->type != TOK_EOF) {
                if (peek(ctx)->type == TOK_COMMA) {
                    advance(ctx);
                    skip_nl(ctx);
                    continue;
                }
                if (peek(ctx)->type == TOK_DOT)
                    advance(ctx);
                spl_tok_t *ftok = advance(ctx);
                char fname[256];
                usize fnl = ftok->len < 255 ? ftok->len : 255;
                memcpy(fname, ftok->lexeme, fnl);
                fname[fnl] = '\0';
                skip_nl(ctx);
                if (peek(ctx)->type == TOK_ASSIGN)
                    advance(ctx);
                skip_nl(ctx);

                spl_expr_result_t fv = spl_parse_expr(ctx, PREC_MIN);
                (void)fv;

                if (strcmp(fname, "ptr") == 0) {
                    spl_emit(ctx, SPL_LADDR, SPL_PTR, offset);
                    spl_emit(ctx, SPL_SWAP, SPL_VOID, 0);
                    spl_emit(ctx, SPL_STORE, SPL_PTR, 0);
                } else if (strcmp(fname, "len") == 0) {
                    spl_emit(ctx, SPL_LADDR, SPL_PTR, offset + (int)sizeof(spl_val_t));
                    spl_emit(ctx, SPL_SWAP, SPL_VOID, 0);
                    spl_emit(ctx, SPL_STORE, SPL_USIZE, 0);
                }
                skip_nl(ctx);
            }
            expect(ctx, TOK_R_BRACE);
        } else if (var_type && (var_type->kind == TYPE_STRUCT || var_type->kind == TYPE_ENUM)) {
            /* Struct/enum initialization */
            usize sz = spl_type_size(var_type);
            if (sz <= sizeof(spl_val_t)) {
                /* Fits in one slot: value on stack, store directly */
                spl_emit(ctx, SPL_LADDR, SPL_PTR, offset);
                spl_emit(ctx, SPL_SWAP, SPL_VOID, 0);
                spl_emit(ctx, SPL_STORE, SPL_PTR, 0);
            } else {
                /* Multi-slot: copy from temp addr to var, slot by slot */
                usize nslots = (sz + sizeof(spl_val_t) - 1) / sizeof(spl_val_t);
                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, offset + (int)(i * sizeof(spl_val_t)));
                    spl_emit(ctx, SPL_SWAP, SPL_VOID, 0);
                    spl_emit(ctx, SPL_STORE, SPL_PTR, 0);
                }
            }
        } else if (var_type && var_type->kind == TYPE_ARRAY) {
            /* Array initialization: store each element in reverse stack order */
            spl_type_info_t *elem = var_type->elem;
            spl_type_t bt = (elem && elem->kind == TYPE_BASIC) ? elem->basic_type : SPL_I32;
            usize stride = spl_type_elem_stride(elem);
            for (int i = (int)var_type->array_len - 1; i >= 0; i--) {
                spl_emit(ctx, SPL_LADDR, SPL_PTR, offset + (int)(i * stride));
                spl_emit(ctx, SPL_SWAP, SPL_VOID, 0);
                spl_emit(ctx, SPL_STORE, bt, 0);
            }
        } else if (var_type && var_type->kind == TYPE_SLICE) {
            /* Slice initialization: stack has [ptr, len] from slice expression.
             * Store len at offset+8, then ptr at offset. */
            spl_emit(ctx, SPL_LADDR, SPL_PTR, offset + (int)sizeof(spl_val_t));
            spl_emit(ctx, SPL_SWAP, SPL_VOID, 0);
            spl_emit(ctx, SPL_STORE, SPL_USIZE, 0);
            spl_emit(ctx, SPL_LADDR, SPL_PTR, offset);
            spl_emit(ctx, SPL_SWAP, SPL_VOID, 0);
            spl_emit(ctx, SPL_STORE, SPL_PTR, 0);
        } else {
            /* Single value store */
            spl_emit(ctx, SPL_LADDR, SPL_PTR, offset);
            spl_type_t bt = (var_type->kind == TYPE_BASIC) ? var_type->basic_type
                            : (var_type->kind == TYPE_PTR) ? SPL_PTR
                                                           : SPL_I32;
            spl_emit(ctx, SPL_SWAP, SPL_VOID, 0);
            spl_emit(ctx, SPL_STORE, bt, 0);
        }
    }

    if (peek(ctx)->type == TOK_SEMICOLON)
        advance(ctx);
}

/* ============================================================
 * Block: { stmt; stmt; ... }
 * ============================================================ */

void spl_parse_block(spl_comp_t *ctx) {
    skip_nl(ctx);
    if (peek(ctx)->type == TOK_L_BRACE) {
        advance(ctx); /* { */
        spl_push_scope(ctx);

        while (peek(ctx)->type != TOK_R_BRACE && peek(ctx)->type != TOK_EOF) {
            spl_parse_stmt(ctx);
            skip_nl(ctx);
        }

        spl_emit_defer_epilogue(ctx, ctx->scope_depth);
        spl_pop_scope(ctx);
        if (peek(ctx)->type == TOK_R_BRACE)
            advance(ctx);
    } else {
        /* Single statement */
        spl_parse_stmt(ctx);
    }
}

/* ============================================================
 * If statement: if expr { ... } [else { ... }]
 * ============================================================ */

static void parse_if_stmt(spl_comp_t *ctx) {
    advance(ctx); /* if */
    skip_nl(ctx);
    spl_expr_result_t cond = spl_parse_expr(ctx, PREC_MIN);
    (void)cond;

    spl_val_t bz_addr = spl_emit_bz(ctx);
    skip_nl(ctx);
    spl_parse_block(ctx);

    spl_val_t jmp_addr = 0;
    skip_nl(ctx);
    if (peek(ctx)->type == KW_ELSE) {
        jmp_addr = spl_emit_jmp(ctx);
        spl_patch_to_here(ctx, bz_addr);
        advance(ctx); /* else */
        skip_nl(ctx);
        spl_parse_block(ctx);
        spl_patch_to_here(ctx, jmp_addr);
    } else {
        spl_patch_to_here(ctx, bz_addr);
    }
}

/* ============================================================
 * While statement: while expr { ... }
 * ============================================================ */

static void parse_while_stmt(spl_comp_t *ctx) {
    spl_val_t loop_start = vec_size(ctx->prog.insns);
    int saved_loop = ctx->in_loop;
    usize saved_continue = ctx->continue_target;
    usize saved_bp_count = ctx->break_patch_count;

    ctx->in_loop = 1;
    ctx->continue_target = loop_start;

    advance(ctx); /* while */
    skip_nl(ctx);
    spl_expr_result_t cond = spl_parse_expr(ctx, PREC_MIN);
    (void)cond;

    spl_val_t bz_addr = spl_emit_bz(ctx);
    skip_nl(ctx);
    spl_parse_block(ctx);

    /* JMP back to loop condition — use relative offset */
    spl_val_t here = vec_size(ctx->prog.insns);
    spl_emit(ctx, SPL_JMP, SPL_VOID, (spl_val_t)((isize)loop_start - (isize)here - 1));
    spl_patch_to_here(ctx, bz_addr);

    /* Patch break statements — compute relative offset */
    for (usize i = saved_bp_count; i < ctx->break_patch_count; i++) {
        spl_patch_to_here(ctx, ctx->break_patches[i]);
    }
    ctx->break_patch_count = saved_bp_count;

    ctx->in_loop = saved_loop;
    ctx->continue_target = saved_continue;
}

/* ============================================================
 * Loop statement: loop { ... }
 * ============================================================ */

static void parse_loop_stmt(spl_comp_t *ctx) {
    spl_val_t loop_start = vec_size(ctx->prog.insns);
    int saved_loop = ctx->in_loop;
    usize saved_continue = ctx->continue_target;
    usize saved_bp_count = ctx->break_patch_count;

    ctx->in_loop = 1;
    ctx->continue_target = loop_start;

    advance(ctx); /* loop */
    skip_nl(ctx);
    spl_parse_block(ctx);

    /* JMP back to loop start — use relative offset */
    spl_val_t here = vec_size(ctx->prog.insns);
    spl_emit(ctx, SPL_JMP, SPL_VOID, (spl_val_t)((isize)loop_start - (isize)here - 1));

    /* Patch break statements — compute relative offset */
    for (usize i = saved_bp_count; i < ctx->break_patch_count; i++) {
        spl_patch_to_here(ctx, ctx->break_patches[i]);
    }
    ctx->break_patch_count = saved_bp_count;

    ctx->in_loop = saved_loop;
    ctx->continue_target = saved_continue;
}

/* ============================================================
 * For loop: for begin..end as i { ... }
 * ============================================================ */

static void parse_for_stmt(spl_comp_t *ctx) {
    advance(ctx); /* for */
    skip_nl(ctx);

    /* Parse the iteration expression(s) */
    spl_expr_result_t start_expr = spl_parse_expr(ctx, PREC_MIN);

    if (peek(ctx)->type == TOK_RANGE) {
        /* ===== Numeric range: for begin..end as i { body } ===== */
        advance(ctx);                  /* .. */
        spl_parse_expr(ctx, PREC_MIN); /* end expression */
        /* Stack: [begin, end] */

        skip_nl(ctx);
        if (peek(ctx)->type != KW_AS) {
            skip_nl(ctx);
            spl_parse_block(ctx);
            return;
        }
        advance(ctx); /* as */

        spl_tok_t *ivar = advance(ctx);
        char iname[256];
        usize inl = ivar->len < 255 ? ivar->len : 255;
        memcpy(iname, ivar->lexeme, inl);
        iname[inl] = '\0';

        spl_push_scope(ctx);
        int ioffset = spl_declare_var(ctx, iname, spl_type_basic(SPL_USIZE), 0);

        /* Stack: [begin, end]; swap so TOS = begin */
        spl_emit(ctx, SPL_SWAP, SPL_VOID, 0);
        /* Store begin to i */
        spl_emit(ctx, SPL_LADDR, SPL_PTR, ioffset);
        spl_emit(ctx, SPL_SWAP, SPL_VOID, 0);
        spl_emit(ctx, SPL_STORE, SPL_USIZE, 0);
        /* Stack: [end] */

        spl_val_t loop_start = vec_size(ctx->prog.insns);

        /* Condition: i < end */
        spl_emit(ctx, SPL_LADDR, SPL_PTR, ioffset);
        spl_emit(ctx, SPL_LOAD, SPL_USIZE, 0);
        spl_emit(ctx, SPL_PICK, SPL_VOID, 1);
        spl_emit(ctx, SPL_ULT, SPL_USIZE, 0);
        spl_val_t bz_addr = spl_emit_bz(ctx);

        int saved_loop = ctx->in_loop;
        usize saved_continue = ctx->continue_target;
        usize saved_bp_count = ctx->break_patch_count;
        ctx->in_loop = 1;
        ctx->continue_target = loop_start;

        skip_nl(ctx);
        spl_parse_block(ctx); /* body */

        /* Increment: i = i + 1 */
        spl_emit(ctx, SPL_LADDR, SPL_PTR, ioffset);
        spl_emit(ctx, SPL_LADDR, SPL_PTR, ioffset);
        spl_emit(ctx, SPL_LOAD, SPL_USIZE, 0);
        spl_emit(ctx, SPL_PUSH, SPL_USIZE, 1);
        spl_emit(ctx, SPL_ADD, SPL_USIZE, 0);
        spl_emit(ctx, SPL_STORE, SPL_USIZE, 0);

        /* JMP back to condition */
        spl_val_t jmp_here = vec_size(ctx->prog.insns);
        spl_emit(ctx, SPL_JMP, SPL_VOID, (spl_val_t)((isize)loop_start - (isize)jmp_here - 1));

        /* Exit: patch bz, drop end */
        spl_patch_to_here(ctx, bz_addr);
        spl_emit(ctx, SPL_DROP, SPL_VOID, 0);

        /* Patch breaks */
        for (usize i = saved_bp_count; i < ctx->break_patch_count; i++)
            spl_patch_to_here(ctx, ctx->break_patches[i]);
        ctx->break_patch_count = saved_bp_count;
        ctx->in_loop = saved_loop;
        ctx->continue_target = saved_continue;

        spl_emit_defer_epilogue(ctx, ctx->scope_depth);
        spl_pop_scope(ctx);
        return;
    }

    /* ===== Slice iteration: for slice [as val |, range as val, idx] ===== */
    if (peek(ctx)->type == TOK_COMMA) {
        advance(ctx);                  /* , */
        spl_parse_expr(ctx, PREC_MIN); /* parse start of range (e.g. 0) */
        if (peek(ctx)->type == TOK_RANGE) {
            advance(ctx); /* .. */
            /* consume optional end expression */
            if (peek(ctx)->type != KW_AS && peek(ctx)->type != TOK_COMMA &&
                peek(ctx)->type != TOK_L_BRACE && peek(ctx)->type != TOK_ENDLINE &&
                peek(ctx)->type != TOK_EOF)
                spl_parse_expr(ctx, PREC_MIN);
        }
        /* Range pushed a value; we manage our own idx, drop it */
        spl_emit(ctx, SPL_DROP, SPL_VOID, 0);
    }

    if (peek(ctx)->type != KW_AS) {
        skip_nl(ctx);
        spl_parse_block(ctx);
        return;
    }
    advance(ctx); /* as */

    /* Parse val variable name */
    spl_tok_t *vtok = advance(ctx);
    char vname[256];
    usize vnl = vtok->len < 255 ? vtok->len : 255;
    memcpy(vname, vtok->lexeme, vnl);
    vname[vnl] = '\0';

    /* Parse optional idx variable name */
    char iname[256] = {0};
    if (peek(ctx)->type == TOK_COMMA) {
        advance(ctx);
        spl_tok_t *itok = advance(ctx);
        usize inl = itok->len < 255 ? itok->len : 255;
        memcpy(iname, itok->lexeme, inl);
        iname[inl] = '\0';
    }

    /* Stack: [slice_addr] or whatever the slice expression left */
    spl_push_scope(ctx);

    int idx_offset = -1;
    if (iname[0])
        idx_offset = spl_declare_var(ctx, iname, spl_type_basic(SPL_USIZE), 0);

    spl_type_info_t *elem_type = NULL;
    if (start_expr.type) {
        if (start_expr.type->kind == TYPE_SLICE)
            elem_type = start_expr.type->elem;
        else if (start_expr.type->kind == TYPE_PTR && start_expr.type->elem)
            elem_type = start_expr.type->elem;
    }
    if (!elem_type)
        elem_type = spl_type_basic(SPL_I32);
    int val_offset = spl_declare_var(ctx, vname, elem_type, 0);

    /* Extract ptr and len from slice, push index=0: stack [ptr, len, idx] */
    spl_emit(ctx, SPL_DUP, SPL_VOID, 0);
    spl_emit(ctx, SPL_LOAD, SPL_PTR, 0);
    spl_emit(ctx, SPL_SWAP, SPL_VOID, 0);
    spl_emit(ctx, SPL_PUSH, SPL_U64, sizeof(spl_val_t));
    spl_emit(ctx, SPL_ADD, SPL_U64, 0);
    spl_emit(ctx, SPL_LOAD, SPL_USIZE, 0);
    spl_emit(ctx, SPL_PUSH, SPL_USIZE, 0);
    /* Stack: [ptr, len, idx=0] */

    spl_val_t loop_start = vec_size(ctx->prog.insns);

    /* Condition: idx < len */
    spl_emit(ctx, SPL_PICK, SPL_VOID, 1); /* copy len */
    spl_emit(ctx, SPL_PICK, SPL_VOID, 1); /* copy idx */
    spl_emit(ctx, SPL_SWAP, SPL_VOID, 0); /* [idx, len] → [len, idx] → SWAP → [idx, len] */
    spl_emit(ctx, SPL_ULT, SPL_USIZE, 0);
    spl_val_t bz_addr = spl_emit_bz(ctx);

    /* Store current idx to idx variable */
    if (idx_offset >= 0) {
        spl_emit(ctx, SPL_PICK, SPL_VOID, 0); /* copy idx (TOS) */
        spl_emit(ctx, SPL_LADDR, SPL_PTR, idx_offset);
        spl_emit(ctx, SPL_SWAP, SPL_VOID, 0);
        spl_emit(ctx, SPL_STORE, SPL_USIZE, 0);
    }

    /* Load slice[idx] and store to val */
    spl_emit(ctx, SPL_PICK, SPL_VOID, 2); /* copy ptr: [ptr, len, idx, ptr] */
    spl_emit(ctx, SPL_PICK, SPL_VOID, 1); /* copy idx: [ptr, len, idx, ptr, idx] */
    usize elem_byte_size = (elem_type && elem_type->byte_size) ? elem_type->byte_size : 4;
    spl_emit(ctx, SPL_PUSH, SPL_U64, elem_byte_size);
    spl_emit(ctx, SPL_MUL, SPL_U64, 0);
    spl_emit(ctx, SPL_ADD, SPL_U64, 0);
    spl_emit(ctx, SPL_LOAD, SPL_I32, 0);
    spl_emit(ctx, SPL_LADDR, SPL_PTR, val_offset);
    spl_emit(ctx, SPL_SWAP, SPL_VOID, 0);
    spl_emit(ctx, SPL_STORE, SPL_I32, 0);

    /* Loop context */
    int saved_loop = ctx->in_loop;
    usize saved_continue = ctx->continue_target;
    usize saved_bp_count = ctx->break_patch_count;
    ctx->in_loop = 1;
    ctx->continue_target = loop_start;

    skip_nl(ctx);
    spl_parse_block(ctx); /* body */

    /* Increment idx */
    spl_emit(ctx, SPL_PICK, SPL_VOID, 0); /* copy idx */
    spl_emit(ctx, SPL_PUSH, SPL_USIZE, 1);
    spl_emit(ctx, SPL_ADD, SPL_USIZE, 0);
    spl_emit(ctx, SPL_SWAP, SPL_VOID, 0);
    spl_emit(ctx, SPL_DROP, SPL_VOID, 0); /* replace old idx with new */

    /* JMP back to condition */
    spl_val_t jmp_here = vec_size(ctx->prog.insns);
    spl_emit(ctx, SPL_JMP, SPL_VOID, (spl_val_t)((isize)loop_start - (isize)jmp_here - 1));

    /* Exit: patch bz, drop idx, len, ptr */
    spl_patch_to_here(ctx, bz_addr);
    spl_emit(ctx, SPL_DROP, SPL_VOID, 0); /* drop idx */
    spl_emit(ctx, SPL_DROP, SPL_VOID, 0); /* drop len */
    spl_emit(ctx, SPL_DROP, SPL_VOID, 0); /* drop ptr */

    /* Patch breaks */
    for (usize i = saved_bp_count; i < ctx->break_patch_count; i++)
        spl_patch_to_here(ctx, ctx->break_patches[i]);
    ctx->break_patch_count = saved_bp_count;
    ctx->in_loop = saved_loop;
    ctx->continue_target = saved_continue;

    spl_pop_scope(ctx);
}

/* ============================================================
 * Break / Continue
 * ============================================================ */

static void parse_break_stmt(spl_comp_t *ctx) {
    advance(ctx); /* break */
    if (!ctx->in_loop) {
        spl_comp_error(ctx, "break outside loop");
    }
    /* Emit JMP with placeholder, add to patch list */
    spl_val_t addr = spl_emit_jmp(ctx);
    if (ctx->break_patch_cap <= ctx->break_patch_count) {
        usize new_cap = ctx->break_patch_cap ? ctx->break_patch_cap * 2 : 8;
        ctx->break_patches = realloc(ctx->break_patches, new_cap * sizeof(usize));
        ctx->break_patch_cap = new_cap;
    }
    ctx->break_patches[ctx->break_patch_count++] = addr;
    if (peek(ctx)->type == TOK_SEMICOLON)
        advance(ctx);
}

static void parse_continue_stmt(spl_comp_t *ctx) {
    advance(ctx); /* continue */
    if (!ctx->in_loop) {
        spl_comp_error(ctx, "continue outside loop");
    }
    /* Emit JMP with relative offset to continue_target */
    spl_val_t here = vec_size(ctx->prog.insns);
    spl_emit(ctx, SPL_JMP, SPL_VOID, (spl_val_t)((isize)ctx->continue_target - (isize)here - 1));
    if (peek(ctx)->type == TOK_SEMICOLON)
        advance(ctx);
}

/* ============================================================
 * Defer statement: defer { ... } or defer expr;
 * ============================================================ */

static void parse_defer_stmt(spl_comp_t *ctx) {
    advance(ctx); /* defer */
    skip_nl(ctx);

    /* Record current position for defer (emits a JMP skip placeholder) */
    spl_emit_defer(ctx);

    /* Parse the deferred statement/block */
    if (peek(ctx)->type == TOK_L_BRACE) {
        spl_parse_block(ctx);
    } else {
        spl_parse_stmt(ctx);
    }

    /* Emit JMP exit placeholder — will be patched at scope exit */
    if (ctx->defer_count > 0) {
        spl_defer_entry_t *e = &ctx->defer_stack[ctx->defer_count - 1];
        e->jmp_exit = spl_emit_jmp(ctx);
    }
}

/* ============================================================
 * Match statement: match expr { .Variant(bindings) => stmt, ... }
 * ============================================================ */

static void parse_match_stmt(spl_comp_t *ctx) {
    advance(ctx); /* match */
    skip_nl(ctx);

    spl_expr_result_t expr = spl_parse_expr(ctx, PREC_MIN);

    /* Determine the enum type (deref pointer if needed) */
    spl_type_info_t *enum_type = expr.type;
    if (enum_type && enum_type->kind == TYPE_PTR && enum_type->elem &&
        enum_type->elem->kind == TYPE_ENUM) {
        enum_type = enum_type->elem;
    }
    if (!enum_type || enum_type->kind != TYPE_ENUM) {
        spl_comp_error(ctx, "match expression must be an enum");
        return;
    }

    /* Save the enum address (pointer value) to a temp slot */
    int addr_offset = ctx->current_local_bytes;
    ctx->current_local_bytes += (int)sizeof(spl_val_t);
    spl_emit(ctx, SPL_LADDR, SPL_PTR, addr_offset);
    spl_emit(ctx, SPL_SWAP, SPL_VOID, 0);
    spl_emit(ctx, SPL_STORE, SPL_PTR, 0);

    skip_nl(ctx);
    if (peek(ctx)->type == TOK_L_BRACE)
        advance(ctx); /* { */

    /* Collect JMP-to-end addresses for patching */
    enum { MAX_MATCH_ARMS = 32 };
    spl_val_t jmp_to_end[MAX_MATCH_ARMS];
    int n_jmps = 0;

    skip_nl(ctx);
    while (peek(ctx)->type != TOK_R_BRACE && peek(ctx)->type != TOK_EOF) {
        skip_nl(ctx);
        if (peek(ctx)->type == TOK_COMMA) {
            advance(ctx);
            skip_nl(ctx);
            continue;
        }

        /* Parse .VariantName */
        if (peek(ctx)->type == TOK_DOT)
            advance(ctx);
        spl_tok_t *vtok = advance(ctx);
        char vname[256];
        usize vnl = vtok->len < 255 ? vtok->len : 255;
        memcpy(vname, vtok->lexeme, vnl);
        vname[vnl] = '\0';

        /* Find the variant in the enum type */
        spl_enum_variant_t *variant = NULL;
        vec_for(enum_type->variants, vi) {
            if (strcmp(vec_at(enum_type->variants, vi).name, vname) == 0) {
                variant = &vec_at(enum_type->variants, vi);
                break;
            }
        }
        if (!variant) {
            spl_comp_error(ctx, "unknown variant '%s' in match", vname);
            break;
        }

        /* Compare tag: load [addr_offset+0] == variant->value */
        /* Load tag and compare */
        spl_emit(ctx, SPL_LADDR, SPL_PTR, addr_offset);
        spl_emit(ctx, SPL_LOAD, SPL_PTR, 0);
        spl_emit(ctx, SPL_PUSH, SPL_USIZE, 0);
        spl_emit(ctx, SPL_ADD, SPL_USIZE, 0);
        spl_emit(ctx, SPL_LOAD, SPL_I32, 0);
        spl_emit(ctx, SPL_PUSH, SPL_I32, variant->value);
        spl_emit(ctx, SPL_EQ, SPL_I32, 0);
        spl_val_t bz_addr = spl_emit_bz(ctx);

        /* This arm matches — parse bindings and body */
        skip_nl(ctx);

        /* Parse binding list: (name, name, ...) */
        int has_parens = 0;
        if (peek(ctx)->type == TOK_L_PAREN) {
            has_parens = 1;
            advance(ctx); /* ( */
            skip_nl(ctx);
        }

        /* Push scope for bindings if there are any */
        int scope_pushed = 0;
        if (has_parens && peek(ctx)->type != TOK_R_PAREN) {
            spl_push_scope(ctx);
            scope_pushed = 1;

            if (variant->data_type && variant->data_type->kind == TYPE_STRUCT) {
                /* Struct data: each struct field is a binding */
                int bi = 0;
                for (;;) {
                    spl_tok_t *btok = advance(ctx);
                    char bname[256];
                    usize bnl = btok->len < 255 ? btok->len : 255;
                    memcpy(bname, btok->lexeme, bnl);
                    bname[bnl] = '\0';

                    spl_type_info_t *btype = spl_type_basic(SPL_I32);
                    usize field_byte_off = 4;
                    if ((usize)bi < vec_size(variant->data_type->fields)) {
                        btype = vec_at(variant->data_type->fields, bi).type;
                        field_byte_off = 4 + vec_at(variant->data_type->fields, bi).offset;
                    }

                    int boffset = spl_declare_var(ctx, bname, btype, 0);

                    /* Load from enum data and store to variable */
                    spl_emit(ctx, SPL_LADDR, SPL_PTR, addr_offset);
                    spl_emit(ctx, SPL_LOAD, SPL_PTR, 0);
                    spl_emit(ctx, SPL_PUSH, SPL_USIZE, field_byte_off);
                    spl_emit(ctx, SPL_ADD, SPL_USIZE, 0);
                    spl_type_t bt = (btype->kind == TYPE_BASIC) ? btype->basic_type
                                    : (btype->kind == TYPE_PTR) ? SPL_PTR
                                                                : SPL_I32;
                    spl_emit(ctx, SPL_LOAD, bt, 0);
                    spl_emit(ctx, SPL_LADDR, SPL_PTR, boffset);
                    spl_emit(ctx, SPL_SWAP, SPL_VOID, 0);
                    spl_emit(ctx, SPL_STORE, bt, 0);

                    bi++;
                    skip_nl(ctx);
                    if (peek(ctx)->type == TOK_COMMA) {
                        advance(ctx);
                        skip_nl(ctx);
                        continue;
                    }
                    break;
                }
            } else if (variant->data_type) {
                /* Simple data type: one binding */
                spl_tok_t *btok = advance(ctx);
                char bname[256];
                usize bnl = btok->len < 255 ? btok->len : 255;
                memcpy(bname, btok->lexeme, bnl);
                bname[bnl] = '\0';

                spl_type_info_t *btype = variant->data_type;
                int boffset = spl_declare_var(ctx, bname, btype, 0);

                /* Load from enum data at offset 4 */
                spl_emit(ctx, SPL_LADDR, SPL_PTR, addr_offset);
                spl_emit(ctx, SPL_LOAD, SPL_PTR, 0);
                spl_emit(ctx, SPL_PUSH, SPL_USIZE, 4);
                spl_emit(ctx, SPL_ADD, SPL_USIZE, 0);
                spl_type_t bt = (btype->kind == TYPE_BASIC) ? btype->basic_type
                                : (btype->kind == TYPE_PTR) ? SPL_PTR
                                                            : SPL_I32;
                spl_emit(ctx, SPL_LOAD, bt, 0);
                spl_emit(ctx, SPL_LADDR, SPL_PTR, boffset);
                spl_emit(ctx, SPL_SWAP, SPL_VOID, 0);
                spl_emit(ctx, SPL_STORE, bt, 0);
            }
        }

        if (has_parens)
            expect(ctx, TOK_R_PAREN);
        skip_nl(ctx);

        /* => (two tokens: = >) */
        if (peek(ctx)->type == TOK_ASSIGN) {
            advance(ctx);
            if (peek(ctx)->type == TOK_GT)
                advance(ctx);
        }
        skip_nl(ctx);

        /* Parse arm body statement */
        spl_parse_stmt(ctx);

        /* Pop scope if we pushed one */
        if (scope_pushed) {
            spl_emit_defer_epilogue(ctx, ctx->scope_depth);
            spl_pop_scope(ctx);
        }

        /* JMP to end (skip remaining arms) */
        if (n_jmps < MAX_MATCH_ARMS)
            jmp_to_end[n_jmps++] = spl_emit_jmp(ctx);

        /* Patch BZ to here (next arm or end) */
        spl_patch_to_here(ctx, bz_addr);

        skip_nl(ctx);
    }

    expect(ctx, TOK_R_BRACE);

    /* Patch all JMPs to end */
    for (int i = 0; i < n_jmps; i++)
        spl_patch_to_here(ctx, jmp_to_end[i]);

    /* Release temp slot */
    ctx->current_local_bytes -= (int)sizeof(spl_val_t);
}

/* ============================================================
 * Extern declaration: #[extern("vm")] fn name(...) ret;
 * ============================================================ */

static void parse_extern_decl(spl_comp_t *ctx) {
    advance(ctx); /* # */
    expect(ctx, TOK_L_BRACKET);
    /* Skip extern("vm") or just look for fn */
    while (peek(ctx)->type != TOK_R_BRACKET && peek(ctx)->type != TOK_EOF)
        advance(ctx);
    if (peek(ctx)->type == TOK_R_BRACKET)
        advance(ctx);

    skip_nl(ctx);

    if (peek(ctx)->type == KW_FN) {
        advance(ctx); /* fn */
        skip_nl(ctx);
        spl_tok_t *fname_tok = advance(ctx);
        char fn_name[256];
        usize fnl = fname_tok->len < 255 ? fname_tok->len : 255;
        memcpy(fn_name, fname_tok->lexeme, fnl);
        fn_name[fnl] = '\0';

        skip_nl(ctx);
        expect(ctx, TOK_L_PAREN);

        /* Count params (we don't store them for extern) */
        int nparams = 0;
        skip_nl(ctx);
        if (peek(ctx)->type != TOK_R_PAREN) {
            while (1) {
                advance(ctx); /* param name */
                skip_nl(ctx);
                if (peek(ctx)->type == TOK_COLON) {
                    advance(ctx); /* : */
                    skip_nl(ctx);
                    spl_parse_type(ctx); /* skip type */
                }
                nparams++;
                skip_nl(ctx);
                if (peek(ctx)->type == TOK_COMMA) {
                    advance(ctx);
                    skip_nl(ctx);
                    continue;
                }
                if (peek(ctx)->type == TOK_ELLIPSIS) {
                    advance(ctx);
                    skip_nl(ctx);
                } /* variadic */
                break;
            }
        }
        expect(ctx, TOK_R_PAREN);
        skip_nl(ctx);

        /* Return type */
        spl_type_info_t *ret_type = spl_type_basic(SPL_VOID);
        if (peek(ctx)->type != TOK_SEMICOLON && peek(ctx)->type != TOK_L_BRACE) {
            ret_type = spl_parse_type(ctx);
            if (!ret_type)
                ret_type = spl_type_basic(SPL_VOID);
            skip_nl(ctx);
        }

        spl_declare_func(ctx, fn_name, ret_type, nparams, 1, 0);

        if (peek(ctx)->type == TOK_SEMICOLON)
            advance(ctx);
    }
}

/* ============================================================
 * Expression statement (may include assignment)
 * ============================================================ */

/* Check if token type is an assignment operator */
static int is_assign_op(spl_tok_type_t t) {
    return t == TOK_ASSIGN || t == TOK_ASSIGN_ADD || t == TOK_ASSIGN_SUB || t == TOK_ASSIGN_MUL ||
           t == TOK_ASSIGN_DIV || t == TOK_ASSIGN_MOD || t == TOK_ASSIGN_AND ||
           t == TOK_ASSIGN_OR || t == TOK_ASSIGN_XOR || t == TOK_ASSIGN_L_SH ||
           t == TOK_ASSIGN_R_SH;
}

static void parse_expr_stmt(spl_comp_t *ctx) {
    skip_nl(ctx);
    if (peek(ctx)->type == KW_RET) {
        parse_ret_stmt(ctx);
        return;
    }

    usize prev = ctx->tok_idx;

    /* Look ahead for assignment operator (scan past expression tokens) */
    int is_assign = 0;
    {
        usize look = ctx->tok_idx;
        while (look < vec_size(ctx->toks)) {
            spl_tok_type_t t = vec_at(ctx->toks, look).type;
            if (t == TOK_SEMICOLON || t == TOK_ENDLINE || t == TOK_L_BRACE || t == TOK_R_BRACE ||
                t == TOK_EOF)
                break;
            if (t == TOK_L_PAREN)
                break; /* function call, not assignment */
            if (is_assign_op(t)) {
                is_assign = 1;
                break;
            }
            look++;
        }
    }

    if (is_assign)
        ctx->addr_of_mode = 1;
    spl_expr_result_t expr = spl_parse_expr(ctx, PREC_MIN);
    if (is_assign)
        ctx->addr_of_mode = 0;

    if (!is_assign && expr.type && expr.type->kind == TYPE_BASIC &&
        expr.type->basic_type != SPL_VOID)
        spl_emit(ctx, SPL_DROP, SPL_VOID, 0);
    if (peek(ctx)->type == TOK_SEMICOLON)
        advance(ctx);
    /* Safety: if no token was consumed, advance to prevent infinite loop */
    if (ctx->tok_idx == prev)
        advance(ctx);
}

/* ============================================================
 * Main statement dispatcher
 * ============================================================ */

void spl_parse_stmt(spl_comp_t *ctx) {
    skip_nl(ctx);

    if (peek(ctx)->type == TOK_EOF || peek(ctx)->type == TOK_R_BRACE)
        return;

    switch (peek(ctx)->type) {
    case KW_RET:
        parse_ret_stmt(ctx);
        break;
    case KW_VAR:
        parse_var_decl(ctx, 0);
        break;
    case KW_CONST:
        parse_var_decl(ctx, 1);
        break;
    case KW_IF:
        parse_if_stmt(ctx);
        break;
    case KW_WHILE:
        parse_while_stmt(ctx);
        break;
    case KW_LOOP:
        parse_loop_stmt(ctx);
        break;
    case KW_FOR:
        parse_for_stmt(ctx);
        break;
    case KW_BREAK:
        parse_break_stmt(ctx);
        break;
    case KW_CONTINUE:
        parse_continue_stmt(ctx);
        break;
    case KW_DEFER:
        parse_defer_stmt(ctx);
        break;
    case KW_MATCH:
        parse_match_stmt(ctx);
        break;
    case KW_TYPE:
        parse_type_decl(ctx);
        break;
    case TOK_L_BRACE:
        spl_parse_block(ctx);
        break;
    case TOK_LINE_COMMENT:
        advance(ctx);
        break;
    case TOK_SHARP: /* #[extern(...)] */
        parse_extern_decl(ctx);
        break;
    default:
        parse_expr_stmt(ctx);
        break;
    }
}