/*
 * par2serial-cc: parser.c - Recursive-descent parser
 * Parses C subset + parallel extensions into AST
 */
#include "parser.h"

/* ── Parser helpers ──────────────────────────────────────── */

#define CUR(p)    ((p)->tokens[(p)->pos])
#define CURTYPE(p) (CUR(p).type)
#define IS(p, t)  (CURTYPE(p) == (t))
#define LOC(p)    (CUR(p).loc)

static Token eat(Parser *p) {
    Token t = CUR(p);
    if (p->pos < p->count - 1) p->pos++;
    return t;
}

static bool check(Parser *p, TokenType t) {
    return CURTYPE(p) == t;
}

static bool match_tok(Parser *p, TokenType t) {
    if (check(p, t)) { eat(p); return true; }
    return false;
}

static Token expect(Parser *p, TokenType t) {
    if (check(p, t)) return eat(p);
    p2s_error(LOC(p), "expected '%s', got '%s' (\"%.*s\")",
              token_type_str(t), token_type_str(CURTYPE(p)),
              (int)CUR(p).text_len, CUR(p).text);
    p->error_count++;
    return CUR(p);
}

static bool is_type_token(Parser *p) {
    TokenType t = CURTYPE(p);
    return t == TOK_VOID || t == TOK_CHAR || t == TOK_SHORT ||
           t == TOK_INT || t == TOK_LONG || t == TOK_FLOAT ||
           t == TOK_DOUBLE || t == TOK_UNSIGNED || t == TOK_SIGNED ||
           t == TOK_STRUCT || t == TOK_UNION || t == TOK_ENUM ||
           t == TOK_CONST || t == TOK_VOLATILE || t == TOK_STATIC ||
           t == TOK_EXTERN || t == TOK_INLINE || t == TOK_REGISTER ||
           t == TOK_RESTRICT || t == TOK_TYPEDEF;
}

/* ── Type parsing ────────────────────────────────────────── */
static ASTType *parse_type(Parser *p);
static ASTNode *parse_expr(Parser *p);
static ASTNode *parse_assign_expr(Parser *p);
static ASTNode *parse_stmt(Parser *p);
static ASTNode *parse_block(Parser *p);
static ASTNode *parse_decl_or_stmt(Parser *p);

static ASTType *parse_base_type(Parser *p) {
    ASTType *ty = ast_new_type(p->arena, TYPE_INT); /* default */
    bool has_type = false;

    /* qualifiers and storage class */
    while (true) {
        if (match_tok(p, TOK_CONST))    { ty->is_const = true; continue; }
        if (match_tok(p, TOK_VOLATILE)) { ty->is_volatile = true; continue; }
        if (match_tok(p, TOK_STATIC))   { ty->is_static = true; continue; }
        if (match_tok(p, TOK_EXTERN))   { continue; }
        if (match_tok(p, TOK_INLINE))   { continue; }
        if (match_tok(p, TOK_REGISTER)) { continue; }
        if (match_tok(p, TOK_RESTRICT)) { ty->is_restrict = true; continue; }
        if (match_tok(p, TOK_TYPEDEF))  { continue; }
        break;
    }

    /* base type */
    if (match_tok(p, TOK_VOID))     { ty->kind = TYPE_VOID; has_type = true; }
    else if (match_tok(p, TOK_CHAR))    { ty->kind = TYPE_CHAR; has_type = true; }
    else if (match_tok(p, TOK_SHORT))   { ty->kind = TYPE_SHORT; has_type = true; }
    else if (match_tok(p, TOK_INT))     { ty->kind = TYPE_INT; has_type = true; }
    else if (match_tok(p, TOK_LONG))    { ty->kind = TYPE_LONG; has_type = true; }
    else if (match_tok(p, TOK_FLOAT))   { ty->kind = TYPE_FLOAT; has_type = true; }
    else if (match_tok(p, TOK_DOUBLE))  { ty->kind = TYPE_DOUBLE; has_type = true; }
    else if (match_tok(p, TOK_UNSIGNED)){
        if (match_tok(p, TOK_CHAR))       ty->kind = TYPE_UNSIGNED_CHAR;
        else if (match_tok(p, TOK_SHORT)) ty->kind = TYPE_UNSIGNED_SHORT;
        else if (match_tok(p, TOK_INT))   ty->kind = TYPE_UNSIGNED_INT;
        else if (match_tok(p, TOK_LONG))  ty->kind = TYPE_UNSIGNED_LONG;
        else                              ty->kind = TYPE_UNSIGNED_INT;
        has_type = true;
    }
    else if (match_tok(p, TOK_SIGNED)) {
        if (match_tok(p, TOK_CHAR))       ty->kind = TYPE_CHAR;
        else if (match_tok(p, TOK_SHORT)) ty->kind = TYPE_SHORT;
        else if (match_tok(p, TOK_INT))   ty->kind = TYPE_INT;
        else if (match_tok(p, TOK_LONG))  ty->kind = TYPE_LONG;
        else                              ty->kind = TYPE_INT;
        has_type = true;
    }
    else if (check(p, TOK_STRUCT) || check(p, TOK_UNION)) {
        bool is_union = check(p, TOK_UNION);
        eat(p);
        ty->kind = is_union ? TYPE_UNION : TYPE_STRUCT;
        if (check(p, TOK_IDENT)) {
            ty->name = arena_strdup(p->arena, CUR(p).text);
            eat(p);
        }
        has_type = true;
        /* skip struct body if present */
        if (check(p, TOK_LBRACE)) {
            eat(p);
            int depth = 1;
            while (depth > 0 && !IS(p, TOK_EOF)) {
                if (check(p, TOK_LBRACE)) depth++;
                if (check(p, TOK_RBRACE)) depth--;
                eat(p);
            }
        }
    }
    else if (check(p, TOK_ENUM)) {
        eat(p);
        ty->kind = TYPE_ENUM;
        if (check(p, TOK_IDENT)) { ty->name = arena_strdup(p->arena, CUR(p).text); eat(p); }
        if (check(p, TOK_LBRACE)) {
            eat(p);
            while (!check(p, TOK_RBRACE) && !IS(p, TOK_EOF)) eat(p);
            expect(p, TOK_RBRACE);
        }
        has_type = true;
    }
    else if (check(p, TOK_IDENT)) {
        /* typedef name - treat as int for now */
        ty->kind = TYPE_INT;
        ty->name = arena_strdup(p->arena, CUR(p).text);
        eat(p);
        has_type = true;
    }

    if (!has_type) ty->kind = TYPE_INT;

    /* trailing long */
    while (match_tok(p, TOK_LONG)) { ty->kind = TYPE_LONG; }

    return ty;
}

static ASTType *parse_type(Parser *p) {
    ASTType *base = parse_base_type(p);
    /* pointer */
    while (check(p, TOK_STAR)) {
        eat(p);
        ASTType *ptr = ast_ptr_type(p->arena, base);
        while (match_tok(p, TOK_CONST)) ptr->is_const = true;
        while (match_tok(p, TOK_RESTRICT)) ptr->is_restrict = true;
        base = ptr;
    }
    return base;
}

/* ── Expression parsing (precedence climbing) ────────────── */

static ASTNode *parse_primary(Parser *p) {
    SourceLoc loc = LOC(p);

    if (check(p, TOK_INT_LIT)) {
        Token t = eat(p);
        ASTNode *n = ast_new(p->arena, NODE_INT_LIT, loc);
        n->int_lit.value = t.int_val;
        return n;
    }
    if (check(p, TOK_FLOAT_LIT)) {
        Token t = eat(p);
        ASTNode *n = ast_new(p->arena, NODE_FLOAT_LIT, loc);
        n->float_lit.value = t.float_val;
        return n;
    }
    if (check(p, TOK_STRING_LIT)) {
        Token t = eat(p);
        ASTNode *n = ast_new(p->arena, NODE_STRING_LIT, loc);
        n->string_lit.value = t.text;
        return n;
    }
    if (check(p, TOK_CHAR_LIT)) {
        Token t = eat(p);
        ASTNode *n = ast_new(p->arena, NODE_CHAR_LIT, loc);
        n->char_lit.value = t.int_val;
        return n;
    }
    if (check(p, TOK_IDENT)) {
        Token t = eat(p);
        ASTNode *n = ast_new(p->arena, NODE_IDENT, loc);
        n->ident.name = t.text;
        return n;
    }
    if (match_tok(p, TOK_LPAREN)) {
        /* check for cast: (type)expr */
        if (is_type_token(p)) {
            ASTType *ty = parse_type(p);
            expect(p, TOK_RPAREN);
            ASTNode *expr = parse_primary(p); /* unary precedence */
            ASTNode *n = ast_new(p->arena, NODE_CAST, loc);
            n->cast.target_type = ty;
            n->cast.expr = expr;
            return n;
        }
        ASTNode *expr = parse_expr(p);
        expect(p, TOK_RPAREN);
        return expr;
    }
    if (check(p, TOK_SIZEOF)) {
        eat(p);
        ASTNode *n = ast_new(p->arena, NODE_SIZEOF, loc);
        if (match_tok(p, TOK_LPAREN)) {
            if (is_type_token(p)) {
                n->size_of.of_type = parse_type(p);
            } else {
                n->size_of.of_expr = parse_expr(p);
            }
            expect(p, TOK_RPAREN);
        } else {
            n->size_of.of_expr = parse_primary(p);
        }
        return n;
    }

    p2s_error(loc, "expected expression, got '%s'", token_type_str(CURTYPE(p)));
    p->error_count++;
    eat(p);
    return ast_new(p->arena, NODE_INT_LIT, loc); /* error recovery */
}

static ASTNode *parse_postfix(Parser *p) {
    ASTNode *expr = parse_primary(p);
    while (true) {
        SourceLoc loc = LOC(p);
        if (match_tok(p, TOK_LBRACKET)) {
            ASTNode *idx = parse_expr(p);
            expect(p, TOK_RBRACKET);
            ASTNode *n = ast_new(p->arena, NODE_INDEX, loc);
            n->index.array = expr;
            n->index.index = idx;
            expr = n;
        }
        else if (match_tok(p, TOK_LPAREN)) {
            /* function call */
            ASTNode *n = ast_new(p->arena, NODE_CALL, loc);
            n->call.func = expr;
            vec_init(&n->call.args);
            if (!check(p, TOK_RPAREN)) {
                vec_push(&n->call.args, parse_assign_expr(p));
                while (match_tok(p, TOK_COMMA))
                    vec_push(&n->call.args, parse_assign_expr(p));
            }
            expect(p, TOK_RPAREN);
            expr = n;
        }
        else if (match_tok(p, TOK_DOT)) {
            Token field = expect(p, TOK_IDENT);
            ASTNode *n = ast_new(p->arena, NODE_MEMBER, loc);
            n->member.object = expr;
            n->member.field = field.text;
            expr = n;
        }
        else if (match_tok(p, TOK_ARROW)) {
            Token field = expect(p, TOK_IDENT);
            ASTNode *n = ast_new(p->arena, NODE_ARROW, loc);
            n->member.object = expr;
            n->member.field = field.text;
            expr = n;
        }
        else if (match_tok(p, TOK_INC)) {
            ASTNode *n = ast_new(p->arena, NODE_POSTFIX, loc);
            n->postfix.op = OP_POST_INC;
            n->postfix.operand = expr;
            expr = n;
        }
        else if (match_tok(p, TOK_DEC)) {
            ASTNode *n = ast_new(p->arena, NODE_POSTFIX, loc);
            n->postfix.op = OP_POST_DEC;
            n->postfix.operand = expr;
            expr = n;
        }
        else break;
    }
    return expr;
}

static ASTNode *parse_unary(Parser *p) {
    SourceLoc loc = LOC(p);

    if (match_tok(p, TOK_INC)) {
        ASTNode *n = ast_new(p->arena, NODE_UNARY, loc);
        n->unary.op = OP_INC;
        n->unary.operand = parse_unary(p);
        return n;
    }
    if (match_tok(p, TOK_DEC)) {
        ASTNode *n = ast_new(p->arena, NODE_UNARY, loc);
        n->unary.op = OP_DEC;
        n->unary.operand = parse_unary(p);
        return n;
    }
    if (match_tok(p, TOK_MINUS)) {
        ASTNode *n = ast_new(p->arena, NODE_UNARY, loc);
        n->unary.op = OP_NEG;
        n->unary.operand = parse_unary(p);
        return n;
    }
    if (match_tok(p, TOK_PLUS)) {
        return parse_unary(p);
    }
    if (match_tok(p, TOK_BANG)) {
        ASTNode *n = ast_new(p->arena, NODE_UNARY, loc);
        n->unary.op = OP_LNOT;
        n->unary.operand = parse_unary(p);
        return n;
    }
    if (match_tok(p, TOK_TILDE)) {
        ASTNode *n = ast_new(p->arena, NODE_UNARY, loc);
        n->unary.op = OP_NOT;
        n->unary.operand = parse_unary(p);
        return n;
    }
    if (match_tok(p, TOK_AMP)) {
        ASTNode *n = ast_new(p->arena, NODE_ADDR_OF, loc);
        n->unary_expr.operand = parse_unary(p);
        return n;
    }
    if (match_tok(p, TOK_STAR)) {
        ASTNode *n = ast_new(p->arena, NODE_DEREF, loc);
        n->unary_expr.operand = parse_unary(p);
        return n;
    }
    return parse_postfix(p);
}

/* Binary expression parsing with precedence */
static int get_precedence(TokenType t) {
    switch (t) {
        case TOK_STAR: case TOK_SLASH: case TOK_PERCENT:  return 13;
        case TOK_PLUS: case TOK_MINUS:                    return 12;
        case TOK_LSHIFT: case TOK_RSHIFT:                 return 11;
        case TOK_LT: case TOK_GT: case TOK_LE: case TOK_GE: return 10;
        case TOK_EQ: case TOK_NEQ:                        return 9;
        case TOK_AMP:                                     return 8;
        case TOK_CARET:                                   return 7;
        case TOK_PIPE:                                    return 6;
        case TOK_AND:                                     return 5;
        case TOK_OR:                                      return 4;
        default: return -1;
    }
}

static OpType token_to_binop(TokenType t) {
    switch (t) {
        case TOK_PLUS:    return OP_ADD;
        case TOK_MINUS:   return OP_SUB;
        case TOK_STAR:    return OP_MUL;
        case TOK_SLASH:   return OP_DIV;
        case TOK_PERCENT: return OP_MOD;
        case TOK_AMP:     return OP_AND;
        case TOK_PIPE:    return OP_OR;
        case TOK_CARET:   return OP_XOR;
        case TOK_LSHIFT:  return OP_LSHIFT;
        case TOK_RSHIFT:  return OP_RSHIFT;
        case TOK_EQ:      return OP_EQ;
        case TOK_NEQ:     return OP_NEQ;
        case TOK_LT:      return OP_LT;
        case TOK_GT:      return OP_GT;
        case TOK_LE:      return OP_LE;
        case TOK_GE:      return OP_GE;
        case TOK_AND:     return OP_LAND;
        case TOK_OR:      return OP_LOR;
        default:          return OP_ADD; /* unreachable */
    }
}

static ASTNode *parse_binary(Parser *p, int min_prec) {
    ASTNode *left = parse_unary(p);
    while (true) {
        int prec = get_precedence(CURTYPE(p));
        if (prec < min_prec) break;
        SourceLoc loc = LOC(p);
        OpType op = token_to_binop(CURTYPE(p));
        eat(p);
        ASTNode *right = parse_binary(p, prec + 1);
        ASTNode *n = ast_new(p->arena, NODE_BINARY, loc);
        n->binary.op = op;
        n->binary.left = left;
        n->binary.right = right;
        left = n;
    }
    return left;
}

static ASTNode *parse_ternary(Parser *p) {
    ASTNode *cond = parse_binary(p, 0);
    if (match_tok(p, TOK_QUESTION)) {
        SourceLoc loc = LOC(p);
        ASTNode *then_expr = parse_expr(p);
        expect(p, TOK_COLON);
        ASTNode *else_expr = parse_ternary(p);
        ASTNode *n = ast_new(p->arena, NODE_TERNARY, loc);
        n->ternary.cond = cond;
        n->ternary.then_expr = then_expr;
        n->ternary.else_expr = else_expr;
        return n;
    }
    return cond;
}

static OpType assign_op(TokenType t) {
    switch (t) {
        case TOK_ASSIGN:     return OP_ASSIGN;
        case TOK_PLUS_EQ:    return OP_ADD_ASSIGN;
        case TOK_MINUS_EQ:   return OP_SUB_ASSIGN;
        case TOK_STAR_EQ:    return OP_MUL_ASSIGN;
        case TOK_SLASH_EQ:   return OP_DIV_ASSIGN;
        case TOK_PERCENT_EQ: return OP_MOD_ASSIGN;
        case TOK_AMP_EQ:     return OP_AND_ASSIGN;
        case TOK_PIPE_EQ:    return OP_OR_ASSIGN;
        case TOK_CARET_EQ:   return OP_XOR_ASSIGN;
        case TOK_LSHIFT_EQ:  return OP_LSHIFT_ASSIGN;
        case TOK_RSHIFT_EQ:  return OP_RSHIFT_ASSIGN;
        default: return (OpType)-1;
    }
}

static bool is_assign_token(TokenType t) {
    return t == TOK_ASSIGN || t == TOK_PLUS_EQ || t == TOK_MINUS_EQ ||
           t == TOK_STAR_EQ || t == TOK_SLASH_EQ || t == TOK_PERCENT_EQ ||
           t == TOK_AMP_EQ || t == TOK_PIPE_EQ || t == TOK_CARET_EQ ||
           t == TOK_LSHIFT_EQ || t == TOK_RSHIFT_EQ;
}

static ASTNode *parse_assign_expr(Parser *p) {
    ASTNode *left = parse_ternary(p);
    if (is_assign_token(CURTYPE(p))) {
        SourceLoc loc = LOC(p);
        OpType op = assign_op(CURTYPE(p));
        eat(p);
        ASTNode *right = parse_assign_expr(p);
        ASTNode *n = ast_new(p->arena, NODE_ASSIGN, loc);
        n->assign.op = op;
        n->assign.target = left;
        n->assign.value = right;
        return n;
    }
    return left;
}

static ASTNode *parse_expr(Parser *p) {
    ASTNode *left = parse_assign_expr(p);
    while (match_tok(p, TOK_COMMA)) {
        SourceLoc loc = LOC(p);
        ASTNode *right = parse_assign_expr(p);
        ASTNode *n = ast_new(p->arena, NODE_COMMA_EXPR, loc);
        n->binary.op = OP_COMMA;
        n->binary.left = left;
        n->binary.right = right;
        left = n;
    }
    return left;
}

/* ── Parallel construct parsing ──────────────────────────── */

static ReduceOp parse_reduce_op(Parser *p) {
    if (check(p, TOK_PLUS))    { eat(p); return REDUCE_ADD; }
    if (check(p, TOK_STAR))    { eat(p); return REDUCE_MUL; }
    if (check(p, TOK_AMP))     { eat(p); return REDUCE_AND; }
    if (check(p, TOK_PIPE))    { eat(p); return REDUCE_OR; }
    if (check(p, TOK_CARET))   { eat(p); return REDUCE_XOR; }
    if (check(p, TOK_IDENT)) {
        if (strcmp(CUR(p).text, "min") == 0) { eat(p); return REDUCE_MIN; }
        if (strcmp(CUR(p).text, "max") == 0) { eat(p); return REDUCE_MAX; }
    }
    p2s_error(LOC(p), "expected reduction operator (+, *, &, |, ^, min, max)");
    p->error_count++;
    return REDUCE_ADD;
}

static ASTNode *parse_parallel_for(Parser *p) {
    SourceLoc loc = LOC(p);
    eat(p); /* parallel_for */
    expect(p, TOK_LPAREN);

    ASTNode *n = ast_new(p->arena, NODE_PARALLEL_FOR, loc);
    Token var = expect(p, TOK_IDENT);
    n->par_for.iter_var = var.text;
    expect(p, TOK_COMMA);
    n->par_for.lo = parse_expr(p);
    expect(p, TOK_COMMA);
    n->par_for.hi = parse_expr(p);
    /* optional step */
    if (match_tok(p, TOK_COMMA)) {
        n->par_for.step = parse_expr(p);
    }
    expect(p, TOK_RPAREN);
    n->par_for.body = parse_block(p);
    return n;
}

static ASTNode *parse_parallel_reduce(Parser *p) {
    SourceLoc loc = LOC(p);
    eat(p); /* parallel_reduce */
    expect(p, TOK_LPAREN);

    ASTNode *n = ast_new(p->arena, NODE_PARALLEL_REDUCE, loc);
    Token acc = expect(p, TOK_IDENT);
    n->par_reduce.accum_var = acc.text;
    expect(p, TOK_COMMA);
    n->par_reduce.op = parse_reduce_op(p);
    expect(p, TOK_COMMA);
    Token ivar = expect(p, TOK_IDENT);
    n->par_reduce.iter_var = ivar.text;
    expect(p, TOK_COMMA);
    n->par_reduce.lo = parse_expr(p);
    expect(p, TOK_COMMA);
    n->par_reduce.hi = parse_expr(p);
    expect(p, TOK_RPAREN);
    n->par_reduce.body = parse_block(p);
    return n;
}

static ASTNode *parse_parallel_scan(Parser *p) {
    SourceLoc loc = LOC(p);
    eat(p); /* parallel_scan */
    expect(p, TOK_LPAREN);

    ASTNode *n = ast_new(p->arena, NODE_PARALLEL_SCAN, loc);
    Token out = expect(p, TOK_IDENT);
    n->par_scan.output_var = out.text;
    expect(p, TOK_COMMA);
    n->par_scan.op = parse_reduce_op(p);
    expect(p, TOK_COMMA);
    Token ivar = expect(p, TOK_IDENT);
    n->par_scan.iter_var = ivar.text;
    expect(p, TOK_COMMA);
    n->par_scan.lo = parse_expr(p);
    expect(p, TOK_COMMA);
    n->par_scan.hi = parse_expr(p);
    expect(p, TOK_RPAREN);
    n->par_scan.body = parse_block(p);
    return n;
}

static ASTNode *parse_parallel_map(Parser *p) {
    SourceLoc loc = LOC(p);
    eat(p); /* parallel_map */
    expect(p, TOK_LPAREN);

    ASTNode *n = ast_new(p->arena, NODE_PARALLEL_MAP, loc);
    Token dst = expect(p, TOK_IDENT);
    n->par_map.dst = dst.text;
    expect(p, TOK_COMMA);
    Token src = expect(p, TOK_IDENT);
    n->par_map.src = src.text;
    expect(p, TOK_COMMA);
    n->par_map.count = parse_expr(p);
    if (match_tok(p, TOK_COMMA)) {
        Token fn = expect(p, TOK_IDENT);
        n->par_map.func_name = fn.text;
    }
    expect(p, TOK_RPAREN);
    n->par_map.body = parse_block(p);
    return n;
}

static ASTNode *parse_tile_hint(Parser *p) {
    SourceLoc loc = LOC(p);
    eat(p); /* tile_hint */
    expect(p, TOK_LPAREN);
    ASTNode *n = ast_new(p->arena, NODE_TILE_HINT, loc);
    Token tgt = expect(p, TOK_IDENT);
    n->tile_hint.target = tgt.text;
    expect(p, TOK_COMMA);
    Token sz = expect(p, TOK_INT_LIT);
    n->tile_hint.tile_size = (int)sz.int_val;
    expect(p, TOK_RPAREN);
    expect(p, TOK_SEMICOLON);
    return n;
}

static ASTNode *parse_simd_hint(Parser *p) {
    SourceLoc loc = LOC(p);
    eat(p); /* simd_hint */
    expect(p, TOK_LPAREN);
    ASTNode *n = ast_new(p->arena, NODE_SIMD_HINT, loc);
    Token tgt = expect(p, TOK_IDENT);
    n->simd_hint.target = tgt.text;
    expect(p, TOK_RPAREN);
    expect(p, TOK_SEMICOLON);
    return n;
}

static ASTNode *parse_memory_layout(Parser *p) {
    SourceLoc loc = LOC(p);
    eat(p); /* memory_layout */
    expect(p, TOK_LPAREN);
    ASTNode *n = ast_new(p->arena, NODE_MEMORY_LAYOUT, loc);
    Token tgt = expect(p, TOK_IDENT);
    n->mem_layout.target = tgt.text;
    expect(p, TOK_COMMA);
    Token layout = expect(p, TOK_IDENT);
    n->mem_layout.soa = (strcmp(layout.text, "SOA") == 0 || strcmp(layout.text, "soa") == 0);
    expect(p, TOK_RPAREN);
    expect(p, TOK_SEMICOLON);
    return n;
}

static ASTNode *parse_barrier(Parser *p) {
    SourceLoc loc = LOC(p);
    eat(p); /* barrier */
    expect(p, TOK_LPAREN);
    expect(p, TOK_RPAREN);
    expect(p, TOK_SEMICOLON);
    return ast_new(p->arena, NODE_BARRIER, loc);
}

/* ── Statement parsing ───────────────────────────────────── */

static ASTNode *parse_block(Parser *p) {
    SourceLoc loc = LOC(p);
    expect(p, TOK_LBRACE);
    ASTNode *n = ast_new(p->arena, NODE_BLOCK, loc);
    vec_init(&n->block.stmts);
    while (!check(p, TOK_RBRACE) && !IS(p, TOK_EOF)) {
        vec_push(&n->block.stmts, parse_decl_or_stmt(p));
    }
    expect(p, TOK_RBRACE);
    return n;
}

static ASTNode *parse_if(Parser *p) {
    SourceLoc loc = LOC(p);
    eat(p); /* if */
    expect(p, TOK_LPAREN);
    ASTNode *cond = parse_expr(p);
    expect(p, TOK_RPAREN);
    ASTNode *n = ast_new(p->arena, NODE_IF, loc);
    n->if_stmt.cond = cond;
    n->if_stmt.then_body = parse_stmt(p);
    if (match_tok(p, TOK_ELSE))
        n->if_stmt.else_body = parse_stmt(p);
    return n;
}

static ASTNode *parse_while(Parser *p) {
    SourceLoc loc = LOC(p);
    eat(p); /* while */
    expect(p, TOK_LPAREN);
    ASTNode *cond = parse_expr(p);
    expect(p, TOK_RPAREN);
    ASTNode *n = ast_new(p->arena, NODE_WHILE, loc);
    n->while_stmt.cond = cond;
    n->while_stmt.body = parse_stmt(p);
    return n;
}

static ASTNode *parse_do_while(Parser *p) {
    SourceLoc loc = LOC(p);
    eat(p); /* do */
    ASTNode *n = ast_new(p->arena, NODE_DO_WHILE, loc);
    n->while_stmt.body = parse_stmt(p);
    expect(p, TOK_WHILE);
    expect(p, TOK_LPAREN);
    n->while_stmt.cond = parse_expr(p);
    expect(p, TOK_RPAREN);
    expect(p, TOK_SEMICOLON);
    return n;
}

static ASTNode *parse_for(Parser *p) {
    SourceLoc loc = LOC(p);
    eat(p); /* for */
    expect(p, TOK_LPAREN);

    ASTNode *n = ast_new(p->arena, NODE_FOR, loc);

    /* init */
    if (!check(p, TOK_SEMICOLON)) {
        if (is_type_token(p)) {
            /* variable declaration in init */
            ASTType *ty = parse_type(p);
            Token name = expect(p, TOK_IDENT);
            ASTNode *decl = ast_new(p->arena, NODE_VAR_DECL, loc);
            decl->var.name = name.text;
            decl->var.var_type = ty;
            if (match_tok(p, TOK_ASSIGN))
                decl->var.init = parse_assign_expr(p);
            n->for_stmt.init = decl;
        } else {
            n->for_stmt.init = parse_expr(p);
        }
    }
    expect(p, TOK_SEMICOLON);

    if (!check(p, TOK_SEMICOLON))
        n->for_stmt.cond = parse_expr(p);
    expect(p, TOK_SEMICOLON);

    if (!check(p, TOK_RPAREN))
        n->for_stmt.step = parse_expr(p);
    expect(p, TOK_RPAREN);

    n->for_stmt.body = parse_stmt(p);
    return n;
}

static ASTNode *parse_return(Parser *p) {
    SourceLoc loc = LOC(p);
    eat(p);
    ASTNode *n = ast_new(p->arena, NODE_RETURN, loc);
    if (!check(p, TOK_SEMICOLON))
        n->ret.value = parse_expr(p);
    expect(p, TOK_SEMICOLON);
    return n;
}

static ASTNode *parse_stmt(Parser *p) {
    switch (CURTYPE(p)) {
    case TOK_LBRACE:          return parse_block(p);
    case TOK_IF:              return parse_if(p);
    case TOK_WHILE:           return parse_while(p);
    case TOK_DO:              return parse_do_while(p);
    case TOK_FOR:             return parse_for(p);
    case TOK_RETURN:          return parse_return(p);
    case TOK_PARALLEL_FOR:    return parse_parallel_for(p);
    case TOK_PARALLEL_REDUCE: return parse_parallel_reduce(p);
    case TOK_PARALLEL_SCAN:   return parse_parallel_scan(p);
    case TOK_PARALLEL_MAP:    return parse_parallel_map(p);
    case TOK_BARRIER:         return parse_barrier(p);
    case TOK_TILE_HINT:       return parse_tile_hint(p);
    case TOK_SIMD_HINT:       return parse_simd_hint(p);
    case TOK_MEMORY_LAYOUT:   return parse_memory_layout(p);
    case TOK_BREAK:
        eat(p); expect(p, TOK_SEMICOLON);
        return ast_new(p->arena, NODE_BREAK, LOC(p));
    case TOK_CONTINUE:
        eat(p); expect(p, TOK_SEMICOLON);
        return ast_new(p->arena, NODE_CONTINUE, LOC(p));
    default: {
        SourceLoc loc = LOC(p);
        ASTNode *expr = parse_expr(p);
        expect(p, TOK_SEMICOLON);
        ASTNode *n = ast_new(p->arena, NODE_EXPR_STMT, loc);
        n->expr_stmt.expr = expr;
        return n;
    }
    }
}

/* ── Declaration parsing ─────────────────────────────────── */

static ASTNode *parse_var_decl(Parser *p, ASTType *type) {
    SourceLoc loc = LOC(p);
    Token name = expect(p, TOK_IDENT);
    ASTNode *n = ast_new(p->arena, NODE_VAR_DECL, loc);
    n->var.name = name.text;
    n->var.var_type = type;

    /* array dimension */
    while (match_tok(p, TOK_LBRACKET)) {
        if (check(p, TOK_RBRACKET)) {
            type = ast_array_type(p->arena, type, -1);
        } else {
            Token sz = expect(p, TOK_INT_LIT);
            type = ast_array_type(p->arena, type, (int)sz.int_val);
        }
        expect(p, TOK_RBRACKET);
        n->var.var_type = type;
    }

    /* initializer */
    if (match_tok(p, TOK_ASSIGN))
        n->var.init = parse_assign_expr(p);

    return n;
}

static ASTNode *parse_func_decl(Parser *p, ASTType *ret_type, const char *name) {
    SourceLoc loc = LOC(p);
    expect(p, TOK_LPAREN);

    ASTNode *n = ast_new(p->arena, NODE_FUNC_DECL, loc);
    n->func.name = name;
    n->func.return_type = ret_type;
    vec_init(&n->func.params);

    if (!check(p, TOK_RPAREN)) {
        if (check(p, TOK_VOID) && p->tokens[p->pos + 1].type == TOK_RPAREN) {
            eat(p); /* void params */
        } else {
            do {
                if (check(p, TOK_ELLIPSIS)) { eat(p); break; }
                ASTType *ptype = parse_type(p);
                ASTNode *param = ast_new(p->arena, NODE_PARAM, LOC(p));
                if (check(p, TOK_IDENT)) {
                    param->param.name = CUR(p).text;
                    eat(p);
                    /* array parameter */
                    while (match_tok(p, TOK_LBRACKET)) {
                        ptype = ast_ptr_type(p->arena, ptype);
                        if (!check(p, TOK_RBRACKET))
                            parse_expr(p); /* skip size */
                        expect(p, TOK_RBRACKET);
                    }
                }
                param->param.param_type = ptype;
                vec_push(&n->func.params, param);
            } while (match_tok(p, TOK_COMMA));
        }
    }
    expect(p, TOK_RPAREN);

    /* function body or just declaration */
    if (check(p, TOK_LBRACE)) {
        n->func.body = parse_block(p);
    } else {
        expect(p, TOK_SEMICOLON);
    }
    return n;
}

static ASTNode *parse_decl_or_stmt(Parser *p) {
    /* preprocessor directives */
    if (check(p, TOK_HASH)) {
        SourceLoc loc = LOC(p);
        Token t = eat(p);
        ASTNode *n = ast_new(p->arena, NODE_PREPROC, loc);
        n->preproc.text = t.text;
        return n;
    }

    /* Parallel constructs */
    switch (CURTYPE(p)) {
    case TOK_PARALLEL_FOR:    return parse_parallel_for(p);
    case TOK_PARALLEL_REDUCE: return parse_parallel_reduce(p);
    case TOK_PARALLEL_SCAN:   return parse_parallel_scan(p);
    case TOK_PARALLEL_MAP:    return parse_parallel_map(p);
    case TOK_BARRIER:         return parse_barrier(p);
    case TOK_TILE_HINT:       return parse_tile_hint(p);
    case TOK_SIMD_HINT:       return parse_simd_hint(p);
    case TOK_MEMORY_LAYOUT:   return parse_memory_layout(p);
    default: break;
    }

    /* Check if this looks like a declaration (type + name) */
    if (is_type_token(p)) {
        ASTType *type = parse_type(p);

        /* struct/union declaration without variable */
        if (check(p, TOK_SEMICOLON)) {
            eat(p);
            ASTNode *n = ast_new(p->arena, NODE_VAR_DECL, LOC(p));
            n->var.var_type = type;
            n->var.name = "";
            return n;
        }

        if (!check(p, TOK_IDENT)) {
            /* might be a statement like `(type)expr;` - fallback */
            return parse_stmt(p);
        }

        Token name = eat(p);

        /* Function declaration? */
        if (check(p, TOK_LPAREN)) {
            /* Look ahead: is this func(params) or func_call() ? */
            /* If preceded by a type, it's a declaration */
            return parse_func_decl(p, type, name.text);
        }

        /* Variable declaration */
        /* Put back to re-parse as var decl is tricky, we handle inline */
        ASTNode *first = ast_new(p->arena, NODE_VAR_DECL, LOC(p));
        first->var.name = name.text;
        first->var.var_type = type;

        /* array */
        while (match_tok(p, TOK_LBRACKET)) {
            if (check(p, TOK_RBRACKET)) {
                type = ast_array_type(p->arena, type, -1);
            } else {
                ASTNode *sz = parse_expr(p);
                int arr_sz = -1;
                if (sz->type == NODE_INT_LIT) arr_sz = (int)sz->int_lit.value;
                type = ast_array_type(p->arena, type, arr_sz);
            }
            expect(p, TOK_RBRACKET);
            first->var.var_type = type;
        }

        if (match_tok(p, TOK_ASSIGN))
            first->var.init = parse_assign_expr(p);

        /* multiple declarations: int a, b, c; */
        /* For simplicity, just handle single for now */
        expect(p, TOK_SEMICOLON);
        return first;
    }

    return parse_stmt(p);
}

/* ── Top-level parser ────────────────────────────────────── */

Parser *parser_create(Arena *arena, Token *tokens, size_t count, const char *filename) {
    Parser *p = (Parser *)arena_alloc(arena, sizeof(Parser));
    p->tokens = tokens;
    p->count = count;
    p->pos = 0;
    p->arena = arena;
    p->error_count = 0;
    p->filename = filename;
    return p;
}

ASTNode *parser_parse(Parser *p) {
    ASTNode *prog = ast_new(p->arena, NODE_PROGRAM, LOC(p));
    vec_init(&prog->program.decls);

    while (!IS(p, TOK_EOF)) {
        ASTNode *decl = parse_decl_or_stmt(p);
        if (decl) vec_push(&prog->program.decls, decl);
        if (p->error_count > 20) {
            p2s_fatal("too many errors, aborting");
            break;
        }
    }
    return prog;
}