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

#include "spl_comp.h"
#include "spl_lex_util.h"
#include <stdlib.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 */
        emit_return(&ctx->emit, &ctx->tctx, ctx->current_ret_type_idx);
    } else {
        spl_expr_result_t val = spl_parse_expr(ctx, PREC_MIN);
        (void)val;
        spl_emit_ret(ctx, ctx->current_ret_type_idx);
    }
    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];
    spl_tok_copy_name(name_tok, vname, sizeof(vname));

    int var_type_idx = -1;
    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_idx = spl_type_parse(&ctx->tctx, 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};
    int inline_lit = 0;
    if (has_init) {
        skip_nl(ctx);
        /* Inline struct/enum/slice literal: var x: Type = { .field = val }
         * Don't call spl_parse_expr — { would be consumed as block expression */
        if (var_type_idx >= 0 && peek(ctx)->type == TOK_L_BRACE &&
            (spl_type_kind(&ctx->tctx, var_type_idx) == TYPE_STRUCT ||
             spl_type_kind(&ctx->tctx, var_type_idx) == TYPE_ENUM ||
             spl_type_kind(&ctx->tctx, var_type_idx) == TYPE_SLICE)) {
            inline_lit = 1;
        } else {
            init = spl_parse_expr(ctx, PREC_MIN);
        }
    }

    if (var_type_idx < 0) {
        var_type_idx = init.type_idx >= 0 ? init.type_idx : spl_type_basic(&ctx->tctx, SPL_I32);
    }
    int offset = spl_declare_var(ctx, vname, var_type_idx, is_const);

    /* Store init value */
    if (has_init) {
        if (inline_lit && (spl_type_kind(&ctx->tctx, var_type_idx) == TYPE_STRUCT ||
                           spl_type_kind(&ctx->tctx, var_type_idx) == TYPE_ENUM ||
                           spl_type_kind(&ctx->tctx, var_type_idx) == TYPE_SLICE)) {
            /* Inline struct/enum/slice literal: var x: Type = { .field = val }
             * Use spl_parse_struct_literal to handle field parsing uniformly. */
            spl_parse_struct_literal(ctx, var_type_idx);
            if (spl_type_kind(&ctx->tctx, var_type_idx) == TYPE_SLICE) {
                /* Slice: copy 2 temp slots to variable */
                emit_frame_copy(&ctx->emit, offset, 2);
            } else {
                spl_emit_store_init(ctx, offset, var_type_idx);
            }
        } else {
            spl_emit_store_init(ctx, offset, var_type_idx);
        }
    }

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

/* Forward declaration for block expr */
static int is_assign_op(spl_tok_type_t t);
static int lookahead_is_assign(spl_comp_t *ctx);

/* ============================================================
 * Block expression: { stmts; [trailing_expr] }
 * Parses a block and returns the trailing expression type.
 * Like Rust/Zig: the last expression without semicolon is the block's value.
 * ============================================================ */

spl_expr_result_t spl_parse_block_expr(spl_comp_t *ctx) {
    advance(ctx); /* { */
    spl_push_scope(ctx);

    spl_expr_result_t result = {0}; /* void by default */

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

        spl_tok_type_t t = peek(ctx)->type;
        int is_keyword =
            (t == KW_RET || t == KW_VAR || t == KW_CONST || t == KW_IF || t == KW_WHILE ||
             t == KW_LOOP || t == KW_FOR || t == KW_BREAK || t == KW_CONTINUE || t == KW_DEFER ||
             t == KW_MATCH || t == KW_TYPE || t == TOK_L_BRACE || t == TOK_AT || t == TOK_SHARP ||
             t == TOK_LINE_COMMENT);

        if (is_keyword) {
            /* Keyword statement — produces void */
            spl_parse_stmt(ctx);
            result = (spl_expr_result_t){0};
        } else {
            /* Expression — look ahead for assignment */
            int is_assign = lookahead_is_assign(ctx);
            int saved_addr = ctx->addr_of_mode;
            if (is_assign)
                ctx->addr_of_mode = 1;
            spl_expr_result_t expr = spl_parse_expr(ctx, PREC_MIN);
            ctx->addr_of_mode = saved_addr;

            skip_nl(ctx);
            if (peek(ctx)->type == TOK_SEMICOLON || peek(ctx)->type == TOK_ENDLINE) {
                /* Expression statement: drop value, consume ; */
                if (!is_assign && expr.type_idx >= 0 &&
                    spl_type_emit_type(&ctx->tctx, expr.type_idx) != SPL_VOID)
                    emit_drop(&ctx->emit);
                while (peek(ctx)->type == TOK_SEMICOLON || peek(ctx)->type == TOK_ENDLINE)
                    advance(ctx);
                result = (spl_expr_result_t){0};
            } else {
                /* Trailing expression — this is the block's value */
                result = expr;
            }
        }
        skip_nl(ctx);
    }

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

/* ============================================================
 * 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 = emit_bz_here(&ctx->emit);
    skip_nl(ctx);
    spl_parse_block(ctx);

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

/* ============================================================
 * Loop context helpers (save/restore loop state for while/loop/for)
 * ============================================================ */

typedef struct {
    int saved_loop;
    usize saved_continue;
    usize saved_bp_count;
    spl_val_t loop_start;
} spl_loop_save_t;

static void loop_enter(spl_comp_t *ctx, spl_loop_save_t *save) {
    save->loop_start = vec_size(ctx->prog.insns);
    save->saved_loop = ctx->in_loop;
    save->saved_continue = ctx->continue_target;
    save->saved_bp_count = ctx->break_patch_count;
    ctx->in_loop = 1;
    ctx->continue_target = save->loop_start;
}

static void loop_exit(spl_comp_t *ctx, spl_loop_save_t *save) {
    spl_val_t here = vec_size(ctx->prog.insns);
    emit_jmp(&ctx->emit, (spl_val_t)((isize)save->loop_start - (isize)here - 1));
    for (usize i = save->saved_bp_count; i < ctx->break_patch_count; i++)
        emit_patch_here(&ctx->emit, ctx->break_patches[i]);
    ctx->break_patch_count = save->saved_bp_count;
    ctx->in_loop = save->saved_loop;
    ctx->continue_target = save->saved_continue;
}

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

static void parse_while_stmt(spl_comp_t *ctx) {
    advance(ctx); /* while */
    skip_nl(ctx);

    spl_loop_save_t save;
    loop_enter(ctx, &save);
    spl_expr_result_t cond = spl_parse_expr(ctx, PREC_MIN);
    (void)cond;
    spl_val_t bz_addr = emit_bz_here(&ctx->emit);
    skip_nl(ctx);
    spl_parse_block(ctx);
    loop_exit(ctx, &save);
    emit_patch_here(&ctx->emit, bz_addr);
}

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

static void parse_loop_stmt(spl_comp_t *ctx) {
    spl_loop_save_t save;
    loop_enter(ctx, &save);
    advance(ctx); /* loop */
    skip_nl(ctx);
    spl_parse_block(ctx);
    loop_exit(ctx, &save);
}

/* ============================================================
 * 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];
        spl_tok_copy_name(ivar, iname, sizeof(iname));

        spl_push_scope(ctx);
        int ioffset = spl_declare_var(ctx, iname, spl_type_basic(&ctx->tctx, SPL_USIZE), 0);

        /* Stack: [begin, end]; swap so TOS = begin */
        emit_swap(&ctx->emit);
        /* Store begin to i */
        emit_store_to_laddr(&ctx->emit, ioffset, SPL_USIZE);
        /* Stack: [end] */

        /* Condition: i < end */
        spl_loop_save_t save;
        loop_enter(ctx, &save);
        emit_laddr(&ctx->emit, ioffset);
        emit_load_usize(&ctx->emit);
        emit_pick(&ctx->emit, 1);
        emit_ult_usize(&ctx->emit);
        spl_val_t bz_addr = emit_bz_here(&ctx->emit);

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

        /* Increment: i = i + 1 */
        emit_laddr(&ctx->emit, ioffset);
        emit_laddr(&ctx->emit, ioffset);
        emit_load_usize(&ctx->emit);
        emit_push_usize(&ctx->emit, 1);
        emit_add_usize(&ctx->emit);
        emit_store_usize(&ctx->emit);

        loop_exit(ctx, &save);

        /* Exit: patch bz, drop end */
        emit_patch_here(&ctx->emit, bz_addr);
        emit_drop(&ctx->emit);

        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 */
        emit_drop(&ctx->emit);
    }

    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];
    spl_tok_copy_name(vtok, vname, sizeof(vname));

    /* Parse optional idx variable name */
    char iname[256] = {0};
    if (peek(ctx)->type == TOK_COMMA) {
        advance(ctx);
        spl_tok_t *itok = advance(ctx);
        spl_tok_copy_name(itok, iname, sizeof(iname));
    }

    /* 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(&ctx->tctx, SPL_USIZE), 0);

    int elem_type_idx = -1;
    if (start_expr.type_idx >= 0) {
        if (spl_type_kind(&ctx->tctx, start_expr.type_idx) == TYPE_SLICE)
            elem_type_idx = spl_type_elem_type(&ctx->tctx, start_expr.type_idx);
        else if (spl_type_kind(&ctx->tctx, start_expr.type_idx) == TYPE_PTR &&
                 spl_type_elem_type(&ctx->tctx, start_expr.type_idx) >= 0)
            elem_type_idx = spl_type_elem_type(&ctx->tctx, start_expr.type_idx);
    }
    if (elem_type_idx < 0)
        elem_type_idx = spl_type_basic(&ctx->tctx, SPL_I32);
    int val_offset = spl_declare_var(ctx, vname, elem_type_idx, 0);

    /* Extract ptr and len from slice, push index=0: stack [ptr, len, idx] */
    emit_dup(&ctx->emit);
    emit_load_ptr(&ctx->emit);
    emit_swap(&ctx->emit);
    emit_ptr_add(&ctx->emit, sizeof(spl_val_t));
    emit_load_usize(&ctx->emit);
    emit_push_usize(&ctx->emit, 0);
    /* Stack: [ptr, len, idx=0] */

    spl_loop_save_t save;
    loop_enter(ctx, &save);

    /* Condition: idx < len */
    emit_pick(&ctx->emit, 1); /* copy len */
    emit_pick(&ctx->emit, 1); /* copy idx */
    emit_swap(&ctx->emit);    /* [idx, len] → [len, idx] → SWAP → [idx, len] */
    emit_ult_usize(&ctx->emit);
    spl_val_t bz_addr = emit_bz_here(&ctx->emit);

    /* Store current idx to idx variable */
    if (idx_offset >= 0) {
        emit_pick(&ctx->emit, 0); /* copy idx (TOS) */
        emit_store_to_laddr(&ctx->emit, idx_offset, SPL_USIZE);
    }

    /* Load slice[idx] and store to val */
    emit_pick(&ctx->emit, 2); /* copy ptr: [ptr, len, idx, ptr] */
    emit_pick(&ctx->emit, 1); /* copy idx: [ptr, len, idx, ptr, idx] */
    usize elem_byte_size = elem_type_idx >= 0 ? spl_type_size(&ctx->tctx, elem_type_idx) : 4;
    emit_push_u64(&ctx->emit, elem_byte_size);
    emit_mul_u64(&ctx->emit);
    emit_add_u64(&ctx->emit);
    if (elem_type_idx >= 0 && !spl_type_is_scalar(&ctx->tctx, elem_type_idx) &&
        spl_type_size(&ctx->tctx, elem_type_idx) > sizeof(spl_val_t)) {
        emit_frame_copy(&ctx->emit, val_offset, spl_type_slot_count(&ctx->tctx, elem_type_idx));
    } else {
        spl_type_t elem_bt =
            elem_type_idx >= 0 ? spl_type_emit_type(&ctx->tctx, elem_type_idx) : SPL_I32;
        emit_load_type(&ctx->emit, elem_bt);
        emit_store_to_laddr(&ctx->emit, val_offset, elem_bt);
    }

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

    /* Increment idx */
    emit_pick(&ctx->emit, 0); /* copy idx */
    emit_push_usize(&ctx->emit, 1);
    emit_add_usize(&ctx->emit);
    emit_swap(&ctx->emit);
    emit_drop(&ctx->emit); /* replace old idx with new */

    loop_exit(ctx, &save);

    /* Exit: patch bz, drop idx, len, ptr */
    emit_patch_here(&ctx->emit, bz_addr);
    emit_drop(&ctx->emit); /* drop idx */
    emit_drop(&ctx->emit); /* drop len */
    emit_drop(&ctx->emit); /* drop ptr */

    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 = emit_jmp_here(&ctx->emit);
    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);
    emit_jmp(&ctx->emit, (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 = emit_jmp_here(&ctx->emit);
    }
}

/* ============================================================
 * Match statement helpers
 * ============================================================ */

/* Parse an enum variant pattern in a match arm: .VariantName
 * Delegates comparison to expr layer (spl_emit_match_enum_cmp).
 * Returns the variant item index, or -1 on error. */
static int parse_match_enum_variant(spl_comp_t *ctx, int enum_type_idx, int val_offset,
                                    int by_value) {
    return spl_emit_match_enum_cmp(ctx, enum_type_idx, val_offset, by_value);
}

/* Parse a value pattern in a match arm (non-enum).
 * Delegates comparison to expr layer (spl_emit_match_value_cmp).
 * Uses PREC_LOGOR internally to prevent => from being consumed as assignment. */
static void parse_match_value_pattern(spl_comp_t *ctx, int val_offset) {
    spl_emit_match_value_cmp(ctx, val_offset);
}

/* Parse enum variant data bindings: (name1, name2, ...)
 * Declares local variables and loads corresponding field data
 * from the matched value's data area (offset 4+).
 * Sets *scope_pushed = 1 if bindings declared. */
static void parse_match_enum_bindings(spl_comp_t *ctx, int enum_type_idx, int variant_item_idx,
                                      int val_offset, int *scope_pushed) {
    advance(ctx); /* ( */
    skip_nl(ctx);
    if (peek(ctx)->type == TOK_R_PAREN) {
        expect(ctx, TOK_R_PAREN);
        return;
    }

    spl_push_scope(ctx);
    *scope_pushed = 1;

    spl_type_item_t *var_item = spl_type_item_at(&ctx->tctx, enum_type_idx, variant_item_idx);
    int data_type_idx = var_item ? var_item->enum_field.type_idx : -1;

    if (data_type_idx >= 0 && spl_type_kind(&ctx->tctx, data_type_idx) == TYPE_STRUCT) {
        /* Multi-field struct destructuring: one binding per struct field */
        spl_type_item_vec_t *data_items = spl_type_items(&ctx->tctx, data_type_idx);
        int bi = 0;
        for (;;) {
            spl_tok_t *btok = advance(ctx);
            char bname[256];
            spl_tok_copy_name(btok, bname, sizeof(bname));

            int btype_idx = spl_type_basic(&ctx->tctx, SPL_I32);
            usize field_byte_off = ENUM_TAG_SIZE; /* skip tag */
            {
                int fcount = 0;
                vec_for(*data_items, di) {
                    spl_type_item_t *fit = &vec_at(*data_items, di);
                    if (fit->item_kind == ITEM_FIELD) {
                        if (fcount == bi) {
                            btype_idx = fit->aggregate_field.type_idx;
                            field_byte_off = ENUM_TAG_SIZE + fit->aggregate_field.offset;
                            break;
                        }
                        fcount++;
                    }
                }
            }
            int boffset = spl_declare_var(ctx, bname, btype_idx, 0);
            emit_load_to_var(&ctx->emit, &ctx->tctx, val_offset, field_byte_off, btype_idx,
                             boffset);

            bi++;
            skip_nl(ctx);
            if (peek(ctx)->type == TOK_COMMA) {
                advance(ctx);
                skip_nl(ctx);
                continue;
            }
            break;
        }
    } else if (data_type_idx >= 0) {
        /* Single-value binding */
        spl_tok_t *btok = advance(ctx);
        char bname[256];
        spl_tok_copy_name(btok, bname, sizeof(bname));

        int boffset = spl_declare_var(ctx, bname, data_type_idx, 0);
        emit_load_to_var(&ctx->emit, &ctx->tctx, val_offset, ENUM_TAG_SIZE, data_type_idx, boffset);
    }

    expect(ctx, TOK_R_PAREN);
}

/* ============================================================
 * Match statement:
 *   Enum:  match expr { .Variant(bindings) => stmt, ... }
 *   Value: match expr { lit, lit => stmt, ..., _ => stmt }
 *
 * Built as enhanced if-else: each arm is a condition chain.
 * Fallthrough (comma-separated patterns) uses BNZ for OR.
 * ============================================================ */

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 match type */
    int is_enum_match = 0;
    int enum_type_idx = -1;
    int type_idx = expr.type_idx;
    if (type_idx >= 0 && spl_type_kind(&ctx->tctx, type_idx) == TYPE_ENUM) {
        is_enum_match = 1;
        enum_type_idx = type_idx;
    } else if (type_idx >= 0 && spl_type_kind(&ctx->tctx, type_idx) == TYPE_PTR &&
               spl_type_elem_type(&ctx->tctx, type_idx) >= 0 &&
               spl_type_kind(&ctx->tctx, spl_type_elem_type(&ctx->tctx, type_idx)) == TYPE_ENUM) {
        is_enum_match = 1;
        enum_type_idx = spl_type_elem_type(&ctx->tctx, type_idx);
    } else if (!(type_idx >= 0 && spl_type_kind(&ctx->tctx, type_idx) == TYPE_BASIC &&
                 spl_type_is_integer(spl_type_basic_type(&ctx->tctx, type_idx)))) {
        spl_comp_error(ctx, "match expression must be an enum or integer type");
        return;
    }

    /* Scalar enums store the raw tag directly; others store a pointer */
    int match_by_value = (type_idx >= 0 && spl_type_kind(&ctx->tctx, type_idx) == TYPE_ENUM &&
                          spl_type_is_scalar(&ctx->tctx, type_idx));

    /* Save match value to temp slot */
    int val_offset = fa_alloc_temp(&ctx->emit.frame, 1);
    emit_store_to_laddr(&ctx->emit, val_offset, spl_type_emit_type(&ctx->tctx, type_idx));

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

    /* Collect JMP-to-end addresses for patching */
    enum { MAX_MATCH_ARMS = 128 };
    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;
        }

        spl_val_t bz_addr = 0;
        int scope_pushed = 0;
        int has_parens = 0;
        spl_val_t body_start = 0;
        spl_val_t bnz_addrs[128];
        int n_bnz = 0;
        int arm_variant_item = -1;

        /* --- Parse arm pattern(s): comma-separated with fallthrough --- */
        if (peek(ctx)->type == KW_ANY) {
            /* _ default arm: always matches, skip comparison */
            advance(ctx);
        } else {
            for (;;) {
                if (is_enum_match) {
                    arm_variant_item =
                        parse_match_enum_variant(ctx, enum_type_idx, val_offset, match_by_value);
                    if (ctx->has_error)
                        break;
                    skip_nl(ctx);
                    if (peek(ctx)->type == TOK_L_PAREN) {
                        has_parens = 1;
                        break; /* bindings → must be last in fallthrough group */
                    }
                } else {
                    parse_match_value_pattern(ctx, val_offset);
                }

                /* Check for more comma-separated patterns */
                skip_nl(ctx);
                if (peek(ctx)->type == TOK_COMMA) {
                    advance(ctx);
                    skip_nl(ctx);
                    if (peek(ctx)->type != TOK_ASSIGN) {
                        if (n_bnz < 128)
                            bnz_addrs[n_bnz++] =
                                emit_bnz_here(&ctx->emit); /* fallthrough to body */
                        continue;
                    }
                    break;
                }
                break;
            }

            /* Last pattern: BZ past body if no match */
            if (!ctx->has_error) {
                bz_addr = emit_bz_here(&ctx->emit);
                body_start = vec_size(ctx->prog.insns);
            }
        }

        /* --- Parse enum bindings (only for last variant) --- */
        if (is_enum_match && has_parens && arm_variant_item >= 0) {
            parse_match_enum_bindings(ctx, enum_type_idx, arm_variant_item, val_offset,
                                      &scope_pushed);
        }

        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 (reuses stmt infrastructure) */
        spl_parse_stmt(ctx);

        /* Pop scope if bindings were declared */
        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++] = emit_jmp_here(&ctx->emit);

        /* Patch BZ to here (next arm or end) */
        if (bz_addr)
            emit_patch_here(&ctx->emit, bz_addr);

        /* Patch BNZ fallthroughs to body start */
        for (int i = 0; i < n_bnz; i++) {
            spl_val_t offset = body_start - bnz_addrs[i] - 1;
            emit_patch(&ctx->emit, bnz_addrs[i], offset);
        }

        skip_nl(ctx);
    }

    expect(ctx, TOK_R_BRACE);

    /* Patch all JMPs to end */
    for (int i = 0; i < n_jmps; i++)
        emit_patch_here(&ctx->emit, jmp_to_end[i]);

    /* Release temp slot */
    fa_free(&ctx->emit.frame, val_offset);
}

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

static void parse_extern_decl(spl_comp_t *ctx) {
    int tt = peek(ctx)->type;
    advance(ctx); /* @ or # */
    skip_nl(ctx);
    if (tt == TOK_AT) {
        /* @extern(vm) fn ... */
        if (peek(ctx)->type == KW_EXTERN) {
            advance(ctx); /* extern */
            skip_nl(ctx);
            if (peek(ctx)->type == TOK_L_PAREN) {
                advance(ctx); /* ( */
                skip_nl(ctx);
                advance(ctx); /* target name */
                skip_nl(ctx);
                if (peek(ctx)->type == TOK_R_PAREN)
                    advance(ctx);
            }
        }
    } else {
        /* #[extern("vm")] fn ... (legacy) */
        expect(ctx, TOK_L_BRACKET);
        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];
        spl_tok_copy_name(fname_tok, fn_name, sizeof(fn_name));

        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_type_parse(&ctx->tctx, 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 */
        int ret_type_idx = spl_type_basic(&ctx->tctx, SPL_VOID);
        if (peek(ctx)->type != TOK_SEMICOLON && peek(ctx)->type != TOK_L_BRACE) {
            ret_type_idx = spl_type_parse(&ctx->tctx, ctx);
            if (ret_type_idx < 0)
                ret_type_idx = spl_type_basic(&ctx->tctx, SPL_VOID);
            skip_nl(ctx);
        }

        spl_declare_func(ctx, fn_name, ret_type_idx, 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;
}

/* Scan ahead to check if the expression at current position is an assignment.
 * Returns 1 if an assignment operator is found before a statement terminator
 * or function-call paren. */
static int lookahead_is_assign(spl_comp_t *ctx) {
    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;
        if (is_assign_op(t))
            return 1;
        look++;
    }
    return 0;
}

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 */
    int is_assign = lookahead_is_assign(ctx);

    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_idx >= 0 &&
        spl_type_emit_type(&ctx->tctx, expr.type_idx) != SPL_VOID)
        emit_drop(&ctx->emit);
    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_AT: { /* @extern(...) or @dbg/@@sizeof expression */
        usize saved = ctx->tok_idx;
        advance(ctx); /* @ */
        skip_nl(ctx);
        spl_tok_t *t = peek(ctx);
        if (t && t->type == KW_EXTERN) {
            ctx->tok_idx = saved;
            parse_extern_decl(ctx);
        } else {
            ctx->tok_idx = saved;
            parse_expr_stmt(ctx);
        }
        break;
    }
    case TOK_SHARP: /* #[extern(...)] (legacy) */
        parse_extern_decl(ctx);
        break;
    default:
        parse_expr_stmt(ctx);
        break;
    }
}