/*
 * par2serial-cc: ir.c - IR construction and AST→IR lowering
 */
#include "ir.h"

/* ── IR Module/Func/Block creation ───────────────────────── */

IRModule *ir_module_create(Arena *arena) {
    IRModule *mod = (IRModule *)arena_alloc(arena, sizeof(IRModule));
    memset(mod, 0, sizeof(IRModule));
    mod->arena = arena;
    vec_init(&mod->preproc_lines);
    vec_init(&mod->global_vars);
    return mod;
}

IRFunc *ir_func_create(IRModule *mod, const char *name) {
    IRFunc *f = (IRFunc *)arena_alloc(mod->arena, sizeof(IRFunc));
    memset(f, 0, sizeof(IRFunc));
    f->name = name;
    vec_init(&f->blocks);
    f->reg_count = 0;
    f->label_count = 0;
    f->next = mod->functions;
    mod->functions = f;
    mod->func_count++;
    return f;
}

IRBlock *ir_block_create(IRFunc *func, const char *name) {
    IRBlock *bb = (IRBlock *)calloc(1, sizeof(IRBlock));
    bb->id = (int)func->blocks.len;
    bb->name = name;
    vec_init(&bb->insts);
    vec_push(&func->blocks, bb);
    return bb;
}

/* ── Value constructors ──────────────────────────────────── */

IRValue ir_reg(int reg, TypeKind type) {
    IRValue v = {0};
    v.kind = VAL_REG;
    v.reg = reg;
    v.type = type;
    return v;
}

IRValue ir_int(int64_t val) {
    IRValue v = {0};
    v.kind = VAL_INT;
    v.int_val = val;
    v.type = TYPE_INT;
    return v;
}

IRValue ir_float(double val) {
    IRValue v = {0};
    v.kind = VAL_FLOAT;
    v.float_val = val;
    v.type = TYPE_FLOAT;
    return v;
}

IRValue ir_label(int id) {
    IRValue v = {0};
    v.kind = VAL_LABEL;
    v.label_id = id;
    return v;
}

IRValue ir_name(const char *name) {
    IRValue v = {0};
    v.kind = VAL_NAME;
    v.name = name;
    return v;
}

IRValue ir_none(void) {
    IRValue v = {0};
    v.kind = VAL_NONE;
    return v;
}

IRValue ir_simd_reg(int reg, TypeKind type, int width) {
    IRValue v = ir_reg(reg, type);
    v.simd_width = width;
    return v;
}

/* ── Instruction emission ────────────────────────────────── */

int ir_new_reg(IRFunc *func) {
    return func->reg_count++;
}

int ir_new_label(IRFunc *func) {
    return func->label_count++;
}

IRInst *ir_emit(Arena *arena, IRBlock *bb, IROp op,
                IRValue dst, IRValue src1, IRValue src2) {
    IRInst *inst = (IRInst *)arena_alloc(arena, sizeof(IRInst));
    memset(inst, 0, sizeof(IRInst));
    inst->op = op;
    inst->dst = dst;
    inst->src1 = src1;
    inst->src2 = src2;
    inst->src3 = ir_none();
    vec_push(&bb->insts, inst);
    return inst;
}

IRInst *ir_emit_comment(Arena *arena, IRBlock *bb, const char *comment) {
    IRInst *inst = (IRInst *)arena_alloc(arena, sizeof(IRInst));
    memset(inst, 0, sizeof(IRInst));
    inst->op = IR_COMMENT;
    inst->comment = comment;
    inst->dst = inst->src1 = inst->src2 = inst->src3 = ir_none();
    vec_push(&bb->insts, inst);
    return inst;
}

/* ══════════════════════════════════════════════════════════
 *  AST → IR Lowering
 * ══════════════════════════════════════════════════════════ */

typedef struct {
    Arena    *arena;
    IRModule *mod;
    IRFunc   *cur_func;
    IRBlock  *cur_block;
    HashMap  *var_regs;   /* variable name → register number */
} IRBuilder;

static IRValue lower_expr(IRBuilder *b, ASTNode *expr);
static void    lower_stmt(IRBuilder *b, ASTNode *stmt);

static int lookup_var(IRBuilder *b, const char *name) {
    void *val = hashmap_get(b->var_regs, name);
    if (val) return (int)(intptr_t)val;
    /* allocate new register for this variable */
    int reg = ir_new_reg(b->cur_func);
    hashmap_put(b->var_regs, name, (void *)(intptr_t)(reg + 1)); /* +1 to avoid NULL */
    return reg;
}

static TypeKind infer_type(ASTNode *expr) {
    if (!expr) return TYPE_INT;
    switch (expr->type) {
    case NODE_FLOAT_LIT: return TYPE_FLOAT;
    case NODE_INT_LIT:   return TYPE_INT;
    default:             return TYPE_INT;
    }
}

static IRValue lower_expr(IRBuilder *b, ASTNode *expr) {
    if (!expr) return ir_int(0);

    switch (expr->type) {
    case NODE_INT_LIT: {
        int r = ir_new_reg(b->cur_func);
        ir_emit(b->arena, b->cur_block, IR_LOAD_IMM,
                ir_reg(r, TYPE_INT), ir_int(expr->int_lit.value), ir_none());
        return ir_reg(r, TYPE_INT);
    }
    case NODE_FLOAT_LIT: {
        int r = ir_new_reg(b->cur_func);
        ir_emit(b->arena, b->cur_block, IR_LOAD_FIMM,
                ir_reg(r, TYPE_FLOAT), ir_float(expr->float_lit.value), ir_none());
        return ir_reg(r, TYPE_FLOAT);
    }
    case NODE_IDENT: {
        int r = lookup_var(b, expr->ident.name);
        return ir_reg(r, TYPE_INT);
    }
    case NODE_STRING_LIT: {
        return ir_name(expr->string_lit.value);
    }
    case NODE_BINARY: {
        IRValue left  = lower_expr(b, expr->binary.left);
        IRValue right = lower_expr(b, expr->binary.right);
        int r = ir_new_reg(b->cur_func);
        bool is_float = (left.type == TYPE_FLOAT || left.type == TYPE_DOUBLE ||
                         right.type == TYPE_FLOAT || right.type == TYPE_DOUBLE);
        IROp op;
        switch (expr->binary.op) {
        case OP_ADD: op = is_float ? IR_FADD : IR_ADD; break;
        case OP_SUB: op = is_float ? IR_FSUB : IR_SUB; break;
        case OP_MUL: op = is_float ? IR_FMUL : IR_MUL; break;
        case OP_DIV: op = is_float ? IR_FDIV : IR_DIV; break;
        case OP_MOD: op = IR_MOD; break;
        case OP_AND: op = IR_AND; break;
        case OP_OR:  op = IR_OR; break;
        case OP_XOR: op = IR_XOR; break;
        case OP_LSHIFT: op = IR_SHL; break;
        case OP_RSHIFT: op = IR_SHR; break;
        case OP_EQ:  op = is_float ? IR_FEQ : IR_EQ; break;
        case OP_NEQ: op = is_float ? IR_FNE : IR_NE; break;
        case OP_LT:  op = is_float ? IR_FLT : IR_LT; break;
        case OP_GT:  op = is_float ? IR_FGT : IR_GT; break;
        case OP_LE:  op = is_float ? IR_FLE : IR_LE; break;
        case OP_GE:  op = is_float ? IR_FGE : IR_GE; break;
        case OP_LAND: op = IR_AND; break;
        case OP_LOR:  op = IR_OR; break;
        default: op = IR_ADD; break;
        }
        TypeKind rtype = is_float ? TYPE_FLOAT : TYPE_INT;
        ir_emit(b->arena, b->cur_block, op,
                ir_reg(r, rtype), left, right);
        return ir_reg(r, rtype);
    }
    case NODE_UNARY: {
        IRValue operand = lower_expr(b, expr->unary.operand);
        int r = ir_new_reg(b->cur_func);
        switch (expr->unary.op) {
        case OP_NEG:
            ir_emit(b->arena, b->cur_block, IR_NEG,
                    ir_reg(r, operand.type), operand, ir_none());
            break;
        case OP_NOT:
            ir_emit(b->arena, b->cur_block, IR_NOT,
                    ir_reg(r, TYPE_INT), operand, ir_none());
            break;
        case OP_LNOT: {
            ir_emit(b->arena, b->cur_block, IR_EQ,
                    ir_reg(r, TYPE_INT), operand, ir_int(0));
            break;
        }
        case OP_INC: {
            ir_emit(b->arena, b->cur_block, IR_ADD,
                    ir_reg(r, operand.type), operand, ir_int(1));
            if (expr->unary.operand->type == NODE_IDENT) {
                int vr = lookup_var(b, expr->unary.operand->ident.name);
                ir_emit(b->arena, b->cur_block, IR_MOVE,
                        ir_reg(vr, operand.type), ir_reg(r, operand.type), ir_none());
            }
            break;
        }
        case OP_DEC: {
            ir_emit(b->arena, b->cur_block, IR_SUB,
                    ir_reg(r, operand.type), operand, ir_int(1));
            if (expr->unary.operand->type == NODE_IDENT) {
                int vr = lookup_var(b, expr->unary.operand->ident.name);
                ir_emit(b->arena, b->cur_block, IR_MOVE,
                        ir_reg(vr, operand.type), ir_reg(r, operand.type), ir_none());
            }
            break;
        }
        default:
            ir_emit(b->arena, b->cur_block, IR_MOVE,
                    ir_reg(r, operand.type), operand, ir_none());
        }
        return ir_reg(r, operand.type);
    }
    case NODE_ASSIGN: {
        IRValue val = lower_expr(b, expr->assign.value);
        if (expr->assign.target->type == NODE_IDENT) {
            int vr = lookup_var(b, expr->assign.target->ident.name);
            if (expr->assign.op == OP_ASSIGN) {
                ir_emit(b->arena, b->cur_block, IR_MOVE,
                        ir_reg(vr, val.type), val, ir_none());
            } else {
                IRValue cur = ir_reg(vr, val.type);
                int tmp = ir_new_reg(b->cur_func);
                IROp op;
                switch (expr->assign.op) {
                case OP_ADD_ASSIGN: op = IR_ADD; break;
                case OP_SUB_ASSIGN: op = IR_SUB; break;
                case OP_MUL_ASSIGN: op = IR_MUL; break;
                case OP_DIV_ASSIGN: op = IR_DIV; break;
                case OP_MOD_ASSIGN: op = IR_MOD; break;
                default: op = IR_ADD; break;
                }
                ir_emit(b->arena, b->cur_block, op,
                        ir_reg(tmp, val.type), cur, val);
                ir_emit(b->arena, b->cur_block, IR_MOVE,
                        ir_reg(vr, val.type), ir_reg(tmp, val.type), ir_none());
            }
            return ir_reg(vr, val.type);
        }
        if (expr->assign.target->type == NODE_INDEX) {
            /* array[idx] = val → STORE */
            IRValue base = lower_expr(b, expr->assign.target->index.array);
            IRValue idx  = lower_expr(b, expr->assign.target->index.index);
            int addr = ir_new_reg(b->cur_func);
            ir_emit(b->arena, b->cur_block, IR_INDEX,
                    ir_reg(addr, TYPE_POINTER), base, idx);
            ir_emit(b->arena, b->cur_block, IR_STORE,
                    ir_reg(addr, TYPE_POINTER), val, ir_none());
            return val;
        }
        return val;
    }
    case NODE_INDEX: {
        IRValue base = lower_expr(b, expr->index.array);
        IRValue idx  = lower_expr(b, expr->index.index);
        int addr = ir_new_reg(b->cur_func);
        ir_emit(b->arena, b->cur_block, IR_INDEX,
                ir_reg(addr, TYPE_POINTER), base, idx);
        int r = ir_new_reg(b->cur_func);
        ir_emit(b->arena, b->cur_block, IR_LOAD,
                ir_reg(r, TYPE_FLOAT), ir_reg(addr, TYPE_POINTER), ir_none());
        return ir_reg(r, TYPE_FLOAT);
    }
    case NODE_CALL: {
        /* Lower arguments */
        IRValue func_val = ir_none();
        if (expr->call.func->type == NODE_IDENT)
            func_val = ir_name(expr->call.func->ident.name);
        /* emit args then call */
        for (size_t i = 0; i < expr->call.args.len; i++)
            lower_expr(b, expr->call.args.data[i]);
        int r = ir_new_reg(b->cur_func);
        ir_emit(b->arena, b->cur_block, IR_CALL,
                ir_reg(r, TYPE_INT), func_val, ir_int((int64_t)expr->call.args.len));
        return ir_reg(r, TYPE_INT);
    }
    case NODE_CAST: {
        IRValue inner = lower_expr(b, expr->cast.expr);
        int r = ir_new_reg(b->cur_func);
        TypeKind target = expr->cast.target_type->kind;
        if ((target == TYPE_FLOAT || target == TYPE_DOUBLE) &&
            inner.type != TYPE_FLOAT && inner.type != TYPE_DOUBLE) {
            ir_emit(b->arena, b->cur_block, IR_INT_TO_FLOAT,
                    ir_reg(r, TYPE_FLOAT), inner, ir_none());
        } else if (inner.type == TYPE_FLOAT && target == TYPE_INT) {
            ir_emit(b->arena, b->cur_block, IR_FLOAT_TO_INT,
                    ir_reg(r, TYPE_INT), inner, ir_none());
        } else {
            ir_emit(b->arena, b->cur_block, IR_MOVE,
                    ir_reg(r, target), inner, ir_none());
        }
        return ir_reg(r, target);
    }
    case NODE_POSTFIX: {
        IRValue operand = lower_expr(b, expr->postfix.operand);
        int old_val = ir_new_reg(b->cur_func);
        ir_emit(b->arena, b->cur_block, IR_MOVE,
                ir_reg(old_val, operand.type), operand, ir_none());
        int new_val = ir_new_reg(b->cur_func);
        IROp op = (expr->postfix.op == OP_POST_INC) ? IR_ADD : IR_SUB;
        ir_emit(b->arena, b->cur_block, op,
                ir_reg(new_val, operand.type), operand, ir_int(1));
        if (expr->postfix.operand->type == NODE_IDENT) {
            int vr = lookup_var(b, expr->postfix.operand->ident.name);
            ir_emit(b->arena, b->cur_block, IR_MOVE,
                    ir_reg(vr, operand.type), ir_reg(new_val, operand.type), ir_none());
        }
        return ir_reg(old_val, operand.type);
    }
    case NODE_TERNARY: {
        IRValue cond_val = lower_expr(b, expr->ternary.cond);
        int lbl_then = ir_new_label(b->cur_func);
        int lbl_else = ir_new_label(b->cur_func);
        int lbl_end  = ir_new_label(b->cur_func);
        int result = ir_new_reg(b->cur_func);

        ir_emit(b->arena, b->cur_block, IR_BRANCH, ir_label(lbl_then), cond_val, ir_none());
        ir_emit(b->arena, b->cur_block, IR_JUMP, ir_label(lbl_else), ir_none(), ir_none());

        /* then */
        ir_emit(b->arena, b->cur_block, IR_LABEL, ir_label(lbl_then), ir_none(), ir_none());
        IRValue then_val = lower_expr(b, expr->ternary.then_expr);
        ir_emit(b->arena, b->cur_block, IR_MOVE,
                ir_reg(result, then_val.type), then_val, ir_none());
        ir_emit(b->arena, b->cur_block, IR_JUMP, ir_label(lbl_end), ir_none(), ir_none());

        /* else */
        ir_emit(b->arena, b->cur_block, IR_LABEL, ir_label(lbl_else), ir_none(), ir_none());
        IRValue else_val = lower_expr(b, expr->ternary.else_expr);
        ir_emit(b->arena, b->cur_block, IR_MOVE,
                ir_reg(result, else_val.type), else_val, ir_none());

        ir_emit(b->arena, b->cur_block, IR_LABEL, ir_label(lbl_end), ir_none(), ir_none());
        return ir_reg(result, then_val.type);
    }
    case NODE_ADDR_OF: {
        if (expr->unary_expr.operand->type == NODE_IDENT) {
            int r = ir_new_reg(b->cur_func);
            ir_emit(b->arena, b->cur_block, IR_LOAD_ADDR,
                    ir_reg(r, TYPE_POINTER),
                    ir_name(expr->unary_expr.operand->ident.name), ir_none());
            return ir_reg(r, TYPE_POINTER);
        }
        return lower_expr(b, expr->unary_expr.operand);
    }
    case NODE_DEREF: {
        IRValue ptr = lower_expr(b, expr->unary_expr.operand);
        int r = ir_new_reg(b->cur_func);
        ir_emit(b->arena, b->cur_block, IR_LOAD,
                ir_reg(r, TYPE_INT), ptr, ir_none());
        return ir_reg(r, TYPE_INT);
    }
    default:
        return ir_int(0);
    }
}

static void lower_stmt(IRBuilder *b, ASTNode *stmt) {
    if (!stmt) return;

    switch (stmt->type) {
    case NODE_BLOCK:
        for (size_t i = 0; i < stmt->block.stmts.len; i++)
            lower_stmt(b, stmt->block.stmts.data[i]);
        break;

    case NODE_EXPR_STMT:
        lower_expr(b, stmt->expr_stmt.expr);
        break;

    case NODE_VAR_DECL: {
        int r = lookup_var(b, stmt->var.name);
        if (stmt->var.init) {
            IRValue val = lower_expr(b, stmt->var.init);
            ir_emit(b->arena, b->cur_block, IR_MOVE,
                    ir_reg(r, val.type), val, ir_none());
        }
        break;
    }

    case NODE_IF: {
        IRValue cond = lower_expr(b, stmt->if_stmt.cond);
        int lbl_then = ir_new_label(b->cur_func);
        int lbl_else = ir_new_label(b->cur_func);
        int lbl_end  = ir_new_label(b->cur_func);

        ir_emit(b->arena, b->cur_block, IR_BRANCH,
                ir_label(lbl_then), cond, ir_none());
        ir_emit(b->arena, b->cur_block, IR_JUMP,
                ir_label(stmt->if_stmt.else_body ? lbl_else : lbl_end),
                ir_none(), ir_none());

        ir_emit(b->arena, b->cur_block, IR_LABEL,
                ir_label(lbl_then), ir_none(), ir_none());
        lower_stmt(b, stmt->if_stmt.then_body);
        ir_emit(b->arena, b->cur_block, IR_JUMP,
                ir_label(lbl_end), ir_none(), ir_none());

        if (stmt->if_stmt.else_body) {
            ir_emit(b->arena, b->cur_block, IR_LABEL,
                    ir_label(lbl_else), ir_none(), ir_none());
            lower_stmt(b, stmt->if_stmt.else_body);
        }

        ir_emit(b->arena, b->cur_block, IR_LABEL,
                ir_label(lbl_end), ir_none(), ir_none());
        break;
    }

    case NODE_WHILE: {
        int lbl_cond = ir_new_label(b->cur_func);
        int lbl_body = ir_new_label(b->cur_func);
        int lbl_end  = ir_new_label(b->cur_func);

        ir_emit(b->arena, b->cur_block, IR_LABEL,
                ir_label(lbl_cond), ir_none(), ir_none());
        IRValue cond = lower_expr(b, stmt->while_stmt.cond);
        ir_emit(b->arena, b->cur_block, IR_BRANCH_FALSE,
                ir_label(lbl_end), cond, ir_none());

        lower_stmt(b, stmt->while_stmt.body);
        ir_emit(b->arena, b->cur_block, IR_JUMP,
                ir_label(lbl_cond), ir_none(), ir_none());

        ir_emit(b->arena, b->cur_block, IR_LABEL,
                ir_label(lbl_end), ir_none(), ir_none());
        break;
    }

    case NODE_FOR: {
        int lbl_cond = ir_new_label(b->cur_func);
        int lbl_body = ir_new_label(b->cur_func);
        int lbl_end  = ir_new_label(b->cur_func);

        if (stmt->for_stmt.init)
            lower_stmt(b, stmt->for_stmt.init);

        ir_emit(b->arena, b->cur_block, IR_LABEL,
                ir_label(lbl_cond), ir_none(), ir_none());

        if (stmt->for_stmt.cond) {
            IRValue cond = lower_expr(b, stmt->for_stmt.cond);
            ir_emit(b->arena, b->cur_block, IR_BRANCH_FALSE,
                    ir_label(lbl_end), cond, ir_none());
        }

        lower_stmt(b, stmt->for_stmt.body);

        if (stmt->for_stmt.step)
            lower_expr(b, stmt->for_stmt.step);

        ir_emit(b->arena, b->cur_block, IR_JUMP,
                ir_label(lbl_cond), ir_none(), ir_none());

        ir_emit(b->arena, b->cur_block, IR_LABEL,
                ir_label(lbl_end), ir_none(), ir_none());
        break;
    }

    case NODE_RETURN: {
        IRValue val = ir_none();
        if (stmt->ret.value) val = lower_expr(b, stmt->ret.value);
        ir_emit(b->arena, b->cur_block, IR_RET, ir_none(), val, ir_none());
        break;
    }

    /* ── Parallel constructs → IR parallel markers ──── */
    case NODE_PARALLEL_FOR: {
        ir_emit_comment(b->arena, b->cur_block, "=== PARALLEL_FOR begin ===");

        int iter_reg = lookup_var(b, stmt->par_for.iter_var);
        IRValue lo = lower_expr(b, stmt->par_for.lo);
        IRValue hi = lower_expr(b, stmt->par_for.hi);

        IRInst *begin = ir_emit(b->arena, b->cur_block, IR_PAR_FOR_BEGIN,
                                ir_reg(iter_reg, TYPE_INT), lo, hi);
        begin->par.iter_var = stmt->par_for.iter_var;
        begin->par.lo_reg = lo.reg;
        begin->par.hi_reg = hi.reg;

        /* Initialize iterator */
        ir_emit(b->arena, b->cur_block, IR_MOVE,
                ir_reg(iter_reg, TYPE_INT), lo, ir_none());

        int lbl_cond = ir_new_label(b->cur_func);
        int lbl_end  = ir_new_label(b->cur_func);

        ir_emit(b->arena, b->cur_block, IR_LABEL,
                ir_label(lbl_cond), ir_none(), ir_none());

        /* cond: iter < hi */
        int cmp_reg = ir_new_reg(b->cur_func);
        ir_emit(b->arena, b->cur_block, IR_LT,
                ir_reg(cmp_reg, TYPE_INT), ir_reg(iter_reg, TYPE_INT), hi);
        ir_emit(b->arena, b->cur_block, IR_BRANCH_FALSE,
                ir_label(lbl_end), ir_reg(cmp_reg, TYPE_INT), ir_none());

        /* body */
        lower_stmt(b, stmt->par_for.body);

        /* step */
        int next = ir_new_reg(b->cur_func);
        ir_emit(b->arena, b->cur_block, IR_ADD,
                ir_reg(next, TYPE_INT), ir_reg(iter_reg, TYPE_INT), ir_int(1));
        ir_emit(b->arena, b->cur_block, IR_MOVE,
                ir_reg(iter_reg, TYPE_INT), ir_reg(next, TYPE_INT), ir_none());

        ir_emit(b->arena, b->cur_block, IR_JUMP,
                ir_label(lbl_cond), ir_none(), ir_none());

        ir_emit(b->arena, b->cur_block, IR_LABEL,
                ir_label(lbl_end), ir_none(), ir_none());

        ir_emit(b->arena, b->cur_block, IR_PAR_FOR_END,
                ir_none(), ir_none(), ir_none());
        ir_emit_comment(b->arena, b->cur_block, "=== PARALLEL_FOR end ===");
        break;
    }

    case NODE_PARALLEL_REDUCE: {
        ir_emit_comment(b->arena, b->cur_block, "=== PARALLEL_REDUCE begin ===");

        int accum_reg = lookup_var(b, stmt->par_reduce.accum_var);
        int iter_reg  = lookup_var(b, stmt->par_reduce.iter_var);
        IRValue lo = lower_expr(b, stmt->par_reduce.lo);
        IRValue hi = lower_expr(b, stmt->par_reduce.hi);

        IRInst *begin = ir_emit(b->arena, b->cur_block, IR_PAR_REDUCE_BEGIN,
                                ir_reg(accum_reg, TYPE_FLOAT), lo, hi);
        begin->par.iter_var = stmt->par_reduce.iter_var;
        begin->par.reduce_op = stmt->par_reduce.op;
        begin->par.accum_reg = accum_reg;

        /* Loop: for(iter = lo; iter < hi; iter++) */
        ir_emit(b->arena, b->cur_block, IR_MOVE,
                ir_reg(iter_reg, TYPE_INT), lo, ir_none());

        int lbl_cond = ir_new_label(b->cur_func);
        int lbl_end  = ir_new_label(b->cur_func);

        ir_emit(b->arena, b->cur_block, IR_LABEL,
                ir_label(lbl_cond), ir_none(), ir_none());

        int cmp_reg = ir_new_reg(b->cur_func);
        ir_emit(b->arena, b->cur_block, IR_LT,
                ir_reg(cmp_reg, TYPE_INT), ir_reg(iter_reg, TYPE_INT), hi);
        ir_emit(b->arena, b->cur_block, IR_BRANCH_FALSE,
                ir_label(lbl_end), ir_reg(cmp_reg, TYPE_INT), ir_none());

        lower_stmt(b, stmt->par_reduce.body);

        int next = ir_new_reg(b->cur_func);
        ir_emit(b->arena, b->cur_block, IR_ADD,
                ir_reg(next, TYPE_INT), ir_reg(iter_reg, TYPE_INT), ir_int(1));
        ir_emit(b->arena, b->cur_block, IR_MOVE,
                ir_reg(iter_reg, TYPE_INT), ir_reg(next, TYPE_INT), ir_none());

        ir_emit(b->arena, b->cur_block, IR_JUMP,
                ir_label(lbl_cond), ir_none(), ir_none());

        ir_emit(b->arena, b->cur_block, IR_LABEL,
                ir_label(lbl_end), ir_none(), ir_none());

        ir_emit(b->arena, b->cur_block, IR_PAR_REDUCE_END,
                ir_reg(accum_reg, TYPE_FLOAT), ir_none(), ir_none());
        ir_emit_comment(b->arena, b->cur_block, "=== PARALLEL_REDUCE end ===");
        break;
    }

    case NODE_PARALLEL_SCAN: {
        ir_emit_comment(b->arena, b->cur_block, "=== PARALLEL_SCAN begin ===");
        /* Lower as sequential scan with markers */
        int iter_reg = lookup_var(b, stmt->par_scan.iter_var);
        IRValue lo = lower_expr(b, stmt->par_scan.lo);
        IRValue hi = lower_expr(b, stmt->par_scan.hi);

        ir_emit(b->arena, b->cur_block, IR_PAR_SCAN_BEGIN,
                ir_none(), lo, hi);

        ir_emit(b->arena, b->cur_block, IR_MOVE,
                ir_reg(iter_reg, TYPE_INT), lo, ir_none());

        int lbl_cond = ir_new_label(b->cur_func);
        int lbl_end  = ir_new_label(b->cur_func);

        ir_emit(b->arena, b->cur_block, IR_LABEL,
                ir_label(lbl_cond), ir_none(), ir_none());
        int cmp = ir_new_reg(b->cur_func);
        ir_emit(b->arena, b->cur_block, IR_LT,
                ir_reg(cmp, TYPE_INT), ir_reg(iter_reg, TYPE_INT), hi);
        ir_emit(b->arena, b->cur_block, IR_BRANCH_FALSE,
                ir_label(lbl_end), ir_reg(cmp, TYPE_INT), ir_none());

        lower_stmt(b, stmt->par_scan.body);

        int next = ir_new_reg(b->cur_func);
        ir_emit(b->arena, b->cur_block, IR_ADD,
                ir_reg(next, TYPE_INT), ir_reg(iter_reg, TYPE_INT), ir_int(1));
        ir_emit(b->arena, b->cur_block, IR_MOVE,
                ir_reg(iter_reg, TYPE_INT), ir_reg(next, TYPE_INT), ir_none());
        ir_emit(b->arena, b->cur_block, IR_JUMP,
                ir_label(lbl_cond), ir_none(), ir_none());

        ir_emit(b->arena, b->cur_block, IR_LABEL,
                ir_label(lbl_end), ir_none(), ir_none());
        ir_emit(b->arena, b->cur_block, IR_PAR_SCAN_END,
                ir_none(), ir_none(), ir_none());
        ir_emit_comment(b->arena, b->cur_block, "=== PARALLEL_SCAN end ===");
        break;
    }

    case NODE_BARRIER:
        ir_emit(b->arena, b->cur_block, IR_BARRIER, ir_none(), ir_none(), ir_none());
        break;

    case NODE_TILE_HINT:
    case NODE_SIMD_HINT:
    case NODE_MEMORY_LAYOUT:
        /* These are hints consumed by optimization passes, stored as comments */
        ir_emit_comment(b->arena, b->cur_block,
                        node_type_str(stmt->type));
        break;

    case NODE_PREPROC:
        /* Pass through */
        break;

    default:
        /* Try as expression */
        lower_expr(b, stmt);
        break;
    }
}

/* ── Top-level lowering ──────────────────────────────────── */

IRModule *ir_build_from_ast(Arena *arena, ASTNode *program) {
    assert(program->type == NODE_PROGRAM);
    IRModule *mod = ir_module_create(arena);

    for (size_t i = 0; i < program->program.decls.len; i++) {
        ASTNode *decl = program->program.decls.data[i];

        if (decl->type == NODE_PREPROC) {
            vec_push(&mod->preproc_lines, decl);
            continue;
        }
        if (decl->type == NODE_VAR_DECL) {
            vec_push(&mod->global_vars, decl);
            continue;
        }
        if (decl->type == NODE_FUNC_DECL) {
            IRFunc *func = ir_func_create(mod, decl->func.name);
            func->param_count = (int)decl->func.params.len;

            IRBlock *entry = ir_block_create(func, "entry");

            IRBuilder builder = {0};
            builder.arena = arena;
            builder.mod = mod;
            builder.cur_func = func;
            builder.cur_block = entry;
            builder.var_regs = hashmap_create();

            /* Map parameters to registers */
            for (size_t j = 0; j < decl->func.params.len; j++) {
                ASTNode *param = decl->func.params.data[j];
                if (param->param.name) {
                    int r = ir_new_reg(func);
                    hashmap_put(builder.var_regs, param->param.name,
                                (void *)(intptr_t)(r + 1));
                }
            }

            if (decl->func.body)
                lower_stmt(&builder, decl->func.body);

            hashmap_destroy(builder.var_regs);
        }
    }

    return mod;
}

/* ══════════════════════════════════════════════════════════
 *  IR Printing
 * ══════════════════════════════════════════════════════════ */

static const char *op_names[] = {
    [IR_ADD]="add", [IR_SUB]="sub", [IR_MUL]="mul", [IR_DIV]="div", [IR_MOD]="mod",
    [IR_NEG]="neg",
    [IR_FADD]="fadd", [IR_FSUB]="fsub", [IR_FMUL]="fmul", [IR_FDIV]="fdiv",
    [IR_AND]="and", [IR_OR]="or", [IR_XOR]="xor", [IR_NOT]="not",
    [IR_SHL]="shl", [IR_SHR]="shr",
    [IR_EQ]="eq", [IR_NE]="ne", [IR_LT]="lt", [IR_GT]="gt", [IR_LE]="le", [IR_GE]="ge",
    [IR_FEQ]="feq", [IR_FNE]="fne", [IR_FLT]="flt", [IR_FGT]="fgt", [IR_FLE]="fle", [IR_FGE]="fge",
    [IR_LOAD]="load", [IR_STORE]="store", [IR_MOVE]="move",
    [IR_LOAD_IMM]="load_imm", [IR_LOAD_FIMM]="load_fimm",
    [IR_LOAD_ADDR]="load_addr", [IR_INDEX]="index",
    [IR_LABEL]="label", [IR_JUMP]="jump", [IR_BRANCH]="branch",
    [IR_BRANCH_FALSE]="branch_false",
    [IR_CALL]="call", [IR_RET]="ret",
    [IR_INT_TO_FLOAT]="i2f", [IR_FLOAT_TO_INT]="f2i",
    [IR_WIDEN]="widen", [IR_NARROW]="narrow",
    [IR_PAR_FOR_BEGIN]="par_for_begin", [IR_PAR_FOR_END]="par_for_end",
    [IR_PAR_REDUCE_BEGIN]="par_reduce_begin", [IR_PAR_REDUCE_END]="par_reduce_end",
    [IR_PAR_SCAN_BEGIN]="par_scan_begin", [IR_PAR_SCAN_END]="par_scan_end",
    [IR_BARRIER]="barrier", [IR_ATOMIC_ADD]="atomic_add",
    [IR_SIMD_LOAD]="simd_load", [IR_SIMD_STORE]="simd_store",
    [IR_SIMD_BROADCAST]="simd_bcast", [IR_SIMD_ADD]="simd_add",
    [IR_SIMD_SUB]="simd_sub", [IR_SIMD_MUL]="simd_mul",
    [IR_SIMD_FMA]="simd_fma", [IR_SIMD_HADD]="simd_hadd",
    [IR_SIMD_MIN]="simd_min", [IR_SIMD_MAX]="simd_max",
    [IR_PREFETCH]="prefetch", [IR_ALLOCA]="alloca",
    [IR_NOP]="nop", [IR_PHI]="phi", [IR_COMMENT]="comment",
};

const char *ir_op_name(IROp op) {
    if (op >= 0 && op < IR_OP_COUNT && op_names[op])
        return op_names[op];
    return "unknown";
}

void ir_print_value(IRValue v) {
    switch (v.kind) {
    case VAL_REG:
        if (v.simd_width > 0) printf("v%d<%d>", v.reg, v.simd_width);
        else printf("%%t%d", v.reg);
        break;
    case VAL_INT:   printf("%ld", (long)v.int_val); break;
    case VAL_FLOAT: printf("%g", v.float_val); break;
    case VAL_LABEL: printf("L%d", v.label_id); break;
    case VAL_NAME:  printf("@%s", v.name); break;
    case VAL_NONE:  break;
    }
}

void ir_print_module(IRModule *mod) {
    printf("=== IR Module (%d functions) ===\n\n", mod->func_count);
    for (IRFunc *f = mod->functions; f; f = f->next) {
        printf("function @%s (%d params, %d regs):\n", f->name, f->param_count, f->reg_count);
        for (size_t i = 0; i < f->blocks.len; i++) {
            IRBlock *bb = f->blocks.data[i];
            printf("  %s:\n", bb->name);
            for (size_t j = 0; j < bb->insts.len; j++) {
                IRInst *inst = bb->insts.data[j];
                if (inst->op == IR_COMMENT) {
                    printf("    ; %s\n", inst->comment);
                    continue;
                }
                if (inst->op == IR_LABEL) {
                    printf("  L%d:\n", inst->dst.label_id);
                    continue;
                }
                printf("    %-18s ", ir_op_name(inst->op));
                if (inst->dst.kind != VAL_NONE) { ir_print_value(inst->dst); printf(" "); }
                if (inst->src1.kind != VAL_NONE) { printf(", "); ir_print_value(inst->src1); }
                if (inst->src2.kind != VAL_NONE) { printf(", "); ir_print_value(inst->src2); }
                if (inst->src3.kind != VAL_NONE) { printf(", "); ir_print_value(inst->src3); }
                printf("\n");
            }
        }
        printf("\n");
    }
}