#include "parse_internal.h"
#include "desugar.h"
#include <stdio.h>
#include <string.h>
#include <stdlib.h>

// === 结构体声明解析 ===
AstNode* parse_struct_decl(Parser* p, ErrorInfo* error) {
    const Token* s_tok = advance(p);  // 跳过 'struct'
    const Token* name = expect(p, TOK_IDENT, error, "struct 后应为结构体名");
    if (!name) return NULL;
    if (!expect(p, TOK_LBRACE, error, "缺少 '{'")) return NULL;

    AstNode* fields[32]; int fcount = 0;
    while (peek(p)->kind != TOK_RBRACE && !error->message) {
        if (fcount >= 32) { error->message = "结构体字段过多 (最多32)"; error->filename = p->filename; error->line = peek(p)->line; error->col = peek(p)->col; return NULL; }
        const Token* fname = expect(p, TOK_IDENT, error, "字段名");
        if (!fname) return NULL;
        if (!expect(p, TOK_COLON, error, "缺少 ':'")) return NULL;
        TypeInfo fti = parse_type_expr(p, error);
        if (fti.kind == TYPE_ERROR) {
            error->filename = p->filename;
            return NULL;
        }
        fields[fcount++] = ast_make_parameter(p->arena,
            arena_strdup_impl(p->arena, fname->start, fname->length),
            fti.kind, fti.struct_name, tok_loc(fname));
        if (peek(p)->kind == TOK_COMMA) advance(p);
        else break;
    }
    if (!expect(p, TOK_RBRACE, error, "缺少 '}'")) return NULL;

    AstNode** farr = arena_alloc_impl(p->arena, fcount * sizeof(AstNode*));
    memcpy(farr, fields, fcount * sizeof(AstNode*));
    return ast_make_struct_decl(p->arena,
        arena_strdup_impl(p->arena, name->start, name->length),
        farr, fcount, tok_loc(s_tok));
}

// === match 语句解析（脱糖为 let + if-else 链）===
// match <expr> { pat1 => { body1 }, pat2 => { body2 }, _ => { body_default } }
//   → { let __match_val = <expr>; if __match_val == pat1 { body1 } else if __match_val == pat2 { body2 } else { body_default } }
static AstNode* parse_match_stmt(Parser* p, ErrorInfo* error) {
    const Token* match_tok = advance(p);  // 跳过 'match'

    // 解析被匹配的表达式
    AstNode* matched = parse_expr(p, error);
    if (!matched) return NULL;

    if (!expect(p, TOK_LBRACE, error, "match 后缺少 '{'")) return NULL;

    // 收集所有分支
    enum { MAX_ARMS = 64 };
    bool   arm_is_wildcard[MAX_ARMS];
    AstNode* arm_pattern[MAX_ARMS];
    AstNode* arm_body[MAX_ARMS];
    int arm_count = 0;

    while (peek(p)->kind != TOK_RBRACE && !error->message) {
        if (arm_count >= MAX_ARMS) {
            error->message = "match 分支过多 (最多64)";
            error->filename = p->filename;
            error->line = peek(p)->line; error->col = peek(p)->col;
            return NULL;
        }
        if (peek(p)->kind == TOK_UNDERSCORE) {
            arm_is_wildcard[arm_count] = true;
            arm_pattern[arm_count] = NULL;
            advance(p);
        } else {
            arm_is_wildcard[arm_count] = false;
            arm_pattern[arm_count] = parse_expr(p, error);
            if (!arm_pattern[arm_count]) return NULL;
        }
        if (!expect(p, TOK_MATCH_ARROW, error, "match 分支缺少 '=>'")) return NULL;
        arm_body[arm_count] = parse_block(p, error);
        if (!arm_body[arm_count]) return NULL;
        arm_count++;
        if (peek(p)->kind == TOK_COMMA) advance(p);
    }
    if (!expect(p, TOK_RBRACE, error, "缺少 '}'")) return NULL;
    if (arm_count == 0) {
        error->message = "match 表达式至少需要一个分支";
        error->filename = p->filename;
        error->line = match_tok->line; error->col = match_tok->col;
        return NULL;
    }
    return desugar_match(p, match_tok, matched, arm_pattern, arm_is_wildcard, arm_body, arm_count);
}

// === 代码块解析 ===
AstNode* parse_block(Parser* p, ErrorInfo* error) {
    if (++parse_depth > MAX_PARSE_DEPTH) {
        error->message = "嵌套过深"; error->filename = p->filename;
        error->line = peek(p)->line; error->col = peek(p)->col;
        parse_depth--; return NULL;
    }
    const Token* open = peek(p);
    if (!expect(p, TOK_LBRACE, error, "缺少 '{'")) { parse_depth--; return NULL; }
    AstNode* stmts[256]; int count = 0;
    while (peek(p)->kind != TOK_RBRACE && peek(p)->kind != TOK_EOF && !error->message) {
        if (count >= 256) { error->message = "代码块语句过多 (最多256)"; error->filename = p->filename; error->line = peek(p)->line; error->col = peek(p)->col; parse_depth--; return NULL; }
        AstNode* s = parse_statement(p, error);
        if (!s) { parse_depth--; return NULL; }
        stmts[count++] = s;
    }
    if (!expect(p, TOK_RBRACE, error, "缺少 '}'")) { parse_depth--; return NULL; }
    AstNode** arr = arena_alloc_impl(p->arena, count * sizeof(AstNode*));
    memcpy(arr, stmts, count * sizeof(AstNode*));
    parse_depth--;
    return ast_make_block(p->arena, arr, count, tok_loc(open));
}

// === 语句解析 ===
AstNode* parse_statement(Parser* p, ErrorInfo* error) {
    const Token* t = peek(p);

    if (t->kind == TOK_LET || t->kind == TOK_VAR) {
        bool is_mut = (advance(p)->kind == TOK_VAR);
        const Token* name = expect(p, TOK_IDENT, error,
            is_mut ? "var 后应为变量名" : "let 后应为变量名");
        if (!name) return NULL;
        // 可选的类型标注
        TypeKind annot_type = TYPE_UNKNOWN;
        bool has_type_annot = false;
        const char* struct_type_name = NULL;
        TypeKind annot_elem_type = 0;
        const char* annot_elem_struct = NULL;
        int64_t annot_arr_size = 0;
        if (match(p, TOK_COLON)) {
            TypeInfo ti = parse_type_expr(p, error);
            if (ti.kind == TYPE_ERROR) return NULL;
            annot_type = ti.kind;
            struct_type_name = ti.struct_name;
            annot_elem_type = ti.element_type;
            annot_elem_struct = ti.element_struct_name;
            annot_arr_size = ti.array_size;
            has_type_annot = true;
        }
        if (!expect(p, TOK_ASSIGN, error, "缺少 '='")) return NULL;
        AstNode* init = parse_expr(p, error);
        if (!init) return NULL;
        if (!expect(p, TOK_SEMICOLON, error, "缺少 ';'")) return NULL;
        return ast_make_let(p->arena,
            arena_strdup_impl(p->arena, name->start, name->length),
            annot_type, has_type_annot, is_mut, init, struct_type_name,
            annot_elem_type, annot_elem_struct, annot_arr_size, tok_loc(t));
    }

    if (t->kind == TOK_IF) {
        // 委托给表达式解析器（含 if let 去糖）
        AstNode* if_expr = parse_expr_prec(p, PREC_NONE, error);
        if (!if_expr) return NULL;
        return if_expr;  // AST_IF_STMT 或 AST_BLOCK(if-let去糖)
    }

    if (t->kind == TOK_WHILE) {
        advance(p);
        AstNode* cond = parse_expr(p, error);
        if (!cond) return NULL;
        AstNode* body = parse_block(p, error);
        if (!body) return NULL;
        return ast_make_while(p->arena, cond, body, tok_loc(t));
    }

    if (t->kind == TOK_FOR) {
        advance(p);
        const Token* var_name = expect(p, TOK_IDENT, error, "for 后应为变量名");
        if (!var_name) return NULL;
        if (!expect(p, TOK_IN, error, "缺少 'in'")) return NULL;
        AstNode* start_expr = parse_expr(p, error);
        if (!start_expr) return NULL;
        if (!expect(p, TOK_TO, error, "缺少 'to'")) return NULL;
        AstNode* end_expr = parse_expr(p, error);
        if (!end_expr) return NULL;
        AstNode* body = parse_block(p, error);
        if (!body) return NULL;
        const char* vname = arena_strdup_impl(p->arena, var_name->start, var_name->length);
        return desugar_for(p, t, vname, start_expr, end_expr, body);
    }

    if (t->kind == TOK_MATCH) {
        return parse_match_stmt(p, error);
    }

    if (t->kind == TOK_GUARD) {
        const Token* guard_tok = advance(p);
        AstNode* cond = parse_expr(p, error);
        if (!cond) return NULL;
        if (!expect(p, TOK_ELSE, error, "guard 缺少 'else'")) return NULL;
        AstNode* body = parse_block(p, error);
        if (!body) return NULL;
        return desugar_guard(p, guard_tok, cond, body);
    }

    if (t->kind == TOK_RETURN) {
        advance(p);
        if (match(p, TOK_SEMICOLON)) {
            return ast_make_return(p->arena, NULL, tok_loc(t));
        }
        AstNode* expr = parse_expr(p, error);
        if (!expr) return NULL;
        if (!expect(p, TOK_SEMICOLON, error, "缺少 ';'")) return NULL;
        return ast_make_return(p->arena, expr, tok_loc(t));
    }

    // 数组元素赋值: ident[expr] = expr ;
    if (t->kind == TOK_IDENT && (t + 1)->kind == TOK_LBRACKET) {
        int ahead_idx = 2;
        int bracket_depth = 1;
        while (bracket_depth > 0 && (t + ahead_idx)->kind != TOK_EOF) {
            if ((t + ahead_idx)->kind == TOK_LBRACKET) bracket_depth++;
            else if ((t + ahead_idx)->kind == TOK_RBRACKET) bracket_depth--;
            if (bracket_depth > 0) ahead_idx++;
        }
        if ((t + ahead_idx + 1)->kind == TOK_ASSIGN) {
            const Token* name = advance(p);  // 消费标识符
            advance(p);  // 消费 '['
            AstNode* index = parse_expr(p, error);
            if (!index) return NULL;
            if (!expect(p, TOK_RBRACKET, error, "缺少 ']'")) return NULL;
            if (!expect(p, TOK_ASSIGN, error, "缺少 '='")) return NULL;
            AstNode* value = parse_expr(p, error);
            if (!value) return NULL;
            if (!expect(p, TOK_SEMICOLON, error, "缺少 ';'")) return NULL;
            return ast_make_array_assign(p->arena,
                arena_strdup_impl(p->arena, name->start, name->length),
                index, value, tok_loc(name));
        }
    }

    // 赋值语句: ident = expr ;
    if (t->kind == TOK_IDENT && (t + 1)->kind == TOK_ASSIGN) {
        const Token* name = advance(p);  // 消费标识符
        advance(p);  // 消费 '='
        AstNode* value = parse_expr(p, error);
        if (!value) return NULL;
        if (!expect(p, TOK_SEMICOLON, error, "缺少 ';'")) return NULL;
        return ast_make_assign(p->arena,
            arena_strdup_impl(p->arena, name->start, name->length),
            value, tok_loc(name));
    }

    // 复合赋值: ident += expr  →  ident = ident + expr
    if (t->kind == TOK_IDENT) {
        TokenKind next_kind = (t + 1)->kind;
        if (next_kind >= TOK_PLUS_EQ && next_kind <= TOK_SLASH_EQ) {
            const Token* name = advance(p);
            TokenKind comp_op = advance(p)->kind;
            BinaryOp binop;
            switch (comp_op) {
                case TOK_PLUS_EQ:  binop = OP_ADD; break;
                case TOK_MINUS_EQ: binop = OP_SUB; break;
                case TOK_STAR_EQ:  binop = OP_MUL; break;
                case TOK_SLASH_EQ: binop = OP_DIV; break;
                default: break;
            }
            AstNode* rhs = parse_expr(p, error);
            if (!rhs) return NULL;
            if (!expect(p, TOK_SEMICOLON, error, "缺少 ';'")) return NULL;
            return desugar_compound_assign(p, name, binop, rhs);
        }
    }

    // 表达式语句
    AstNode* expr = parse_expr(p, error);
    if (!expr) return NULL;
    if (!expect(p, TOK_SEMICOLON, error, "缺少 ';'")) return NULL;
    return ast_make_expr_stmt(p->arena, expr, tok_loc(t));
}

// === 函数解析 ===
AstNode* parse_function(Parser* p, bool is_pub, ErrorInfo* error) {
    const Token* fn_tok = advance(p);  // fn
    const Token* name = expect(p, TOK_IDENT, error, "fn 后应为函数名");
    if (!name) return NULL;
    // 泛型类型参数: <T, U, ...>
    const char* type_params[8]; int tp_count = 0;
    if (peek(p)->kind == TOK_LT) {
        advance(p);  // 跳过 '<'
        while (peek(p)->kind != TOK_GT && !error->message) {
            if (tp_count >= 8) { error->message = "类型参数过多 (最多8)"; error->filename = p->filename; error->line = peek(p)->line; error->col = peek(p)->col; return NULL; }
            const Token* tp = expect(p, TOK_IDENT, error, "类型参数名");
            if (!tp) return NULL;
            type_params[tp_count++] = arena_strdup_impl(p->arena, tp->start, tp->length);
            if (peek(p)->kind == TOK_COMMA) advance(p); else break;
        }
        if (!expect(p, TOK_GT, error, "缺少 '>'")) return NULL;
    }
    if (!expect(p, TOK_LPAREN, error, "缺少 '('")) return NULL;

    // 参数列表
    AstNode* params[64]; int pcount = 0;
    while (peek(p)->kind != TOK_RPAREN && !error->message) {
        if (pcount >= 64) { error->message = "函数参数过多 (最多64)"; error->filename = p->filename; error->line = peek(p)->line; error->col = peek(p)->col; return NULL; }
        const Token* pname = expect(p, TOK_IDENT, error, "参数名");
        if (!pname) return NULL;
        if (!expect(p, TOK_COLON, error, "缺少 ':'")) return NULL;
        TypeInfo pti = parse_type_expr(p, error);
        if (pti.kind == TYPE_ERROR) return NULL;
        params[pcount++] = ast_make_parameter(p->arena,
            arena_strdup_impl(p->arena, pname->start, pname->length),
            pti.kind, pti.struct_name, tok_loc(pname));
        if (match(p, TOK_COMMA)) continue;
        else break;
    }
    if (!expect(p, TOK_RPAREN, error, "缺少 ')'")) return NULL;

    // 返回类型
    TypeKind ret = TYPE_VOID;
    const char* ret_struct_name = NULL;
    if (match(p, TOK_ARROW)) {
        TypeInfo rti = parse_type_expr(p, error);
        if (rti.kind == TYPE_ERROR) return NULL;
        ret = rti.kind;
        ret_struct_name = rti.struct_name;
    }

    // trait 方法签名或普通函数体
    AstNode* body = NULL;
    if (match(p, TOK_SEMICOLON)) {
        body = NULL;  // trait 方法签名，无实现
    } else {
        body = parse_block(p, error);
        if (!body) return NULL;
    }

    AstNode** parr = arena_alloc_impl(p->arena, pcount * sizeof(AstNode*));
    memcpy(parr, params, pcount * sizeof(AstNode*));
    const char** tparr = NULL;
    if (tp_count > 0) {
        tparr = arena_alloc_impl(p->arena, tp_count * sizeof(const char*));
        memcpy(tparr, type_params, tp_count * sizeof(const char*));
    }
    return ast_make_function(p->arena,
        arena_strdup_impl(p->arena, name->start, name->length),
        parr, pcount, ret, ret_struct_name, body, is_pub, tparr, tp_count, tok_loc(fn_tok));
}

// === 模块文件加载 ===
static AstNode* load_module(Arena* a, const char* parent_file,
                             const char* mod_name, ErrorInfo* error) {
    // 构造模块文件路径: 同目录下 mod_name.l
    char mod_path[512];
    const char* last_slash = strrchr(parent_file, '/');
    const char* last_bs = strrchr(parent_file, '\\');
    const char* dir_end = parent_file;
    if (last_slash && last_slash > dir_end) dir_end = last_slash;
    if (last_bs && last_bs > dir_end) dir_end = last_bs;
    if (dir_end != parent_file) {
        size_t dir_len = dir_end - parent_file + 1;
        memcpy(mod_path, parent_file, dir_len);
        snprintf(mod_path + dir_len, sizeof(mod_path) - dir_len, "%s.l", mod_name);
    } else {
        snprintf(mod_path, sizeof(mod_path), "%s.l", mod_name);
    }
    // 读取文件
    FILE* f = fopen(mod_path, "rb");
    if (!f) {
        error->message = "无法打开模块文件"; error->filename = mod_path;
        error->line = 0; error->col = 0;
        return NULL;
    }
    fseek(f, 0, SEEK_END);
    size_t sz = ftell(f); fseek(f, 0, SEEK_SET);
    char* src = arena_alloc_impl(a, sz + 1);
    if (!src) { fclose(f); return NULL; }
    fread(src, 1, sz, f); src[sz] = '\0'; fclose(f);

    size_t tc;
    ErrorInfo lex_err = {0};
    Token* toks = lex(a, src, mod_path, &tc, &lex_err);
    if (!toks) { *error = lex_err; return NULL; }

    AstNode* ast = parse(a, toks, tc, mod_path, error);
    return ast;
}

// === 程序入口 ===
AstNode* parse(Arena* a, const Token* tokens, size_t count,
               const char* filename, ErrorInfo* error) {
    Parser p = {.tokens = tokens, .count = count, .pos = 0,
                .filename = filename, .arena = a};
    AstNode* functions[256]; int fn_count = 0;
    AstNode* structs[64];   int struct_count = 0;
    AstNode* aliases[64];   int alias_count = 0;
    AstNode* enums[64];     int enum_count = 0;
    AstNode* impls[64];     int impl_count = 0;
    while (peek(&p)->kind != TOK_EOF && !error->message) {
        // pub 前缀
        bool is_pub = false;
        if (peek(&p)->kind == TOK_PUB) {
            is_pub = true; advance(&p);
        }
        if (peek(&p)->kind == TOK_TRAIT) {
            const Token* tt = advance(&p);
            const Token* tname = expect(&p, TOK_IDENT, error, "trait 后应为接口名");
            if (!tname) return NULL;
            if (!expect(&p, TOK_LBRACE, error, "缺少 '{'")) return NULL;
            AstNode* methods[64]; int mcount = 0;
            while (peek(&p)->kind != TOK_RBRACE && !error->message) {
                if (mcount >= 64) { error->message = "trait 方法过多(最多64)"; error->filename = p.filename; error->line = peek(&p)->line; error->col = peek(&p)->col; return NULL; }
                if (peek(&p)->kind != TOK_FN) { error->message = "trait 内只允许 fn"; error->filename = p.filename; error->line = peek(&p)->line; error->col = peek(&p)->col; return NULL; }
                // trait 方法只解析签名（body 为空）
                AstNode* m = parse_function(&p, false, error);
                if (!m) return NULL;
                m->as.function.body = NULL;  // trait 方法无实现
                methods[mcount++] = m;
                if (peek(&p)->kind == TOK_COMMA) advance(&p);
            }
            if (!expect(&p, TOK_RBRACE, error, "缺少 '}'")) return NULL;
            AstNode** marr = arena_alloc_impl(p.arena, mcount * sizeof(AstNode*));
            memcpy(marr, methods, mcount * sizeof(AstNode*));
            if (impl_count >= 64) { error->message = "trait 过多(最多64)"; error->filename = p.filename; error->line = peek(&p)->line; error->col = peek(&p)->col; return NULL; }
            impls[impl_count++] = ast_make_trait_decl(p.arena,
                arena_strdup_impl(p.arena, tname->start, tname->length),
                marr, mcount, tok_loc(tt));
        } else if (peek(&p)->kind == TOK_STRUCT) {
            if (struct_count >= 64) { error->message = "结构体过多 (最多64)"; error->filename = p.filename; error->line = peek(&p)->line; error->col = peek(&p)->col; return NULL; }
            structs[struct_count++] = parse_struct_decl(&p, error);
        } else if (peek(&p)->kind == TOK_TYPE) {
            if (alias_count >= 64) { error->message = "类型别名过多 (最多64)"; error->filename = p.filename; error->line = peek(&p)->line; error->col = peek(&p)->col; return NULL; }
            const Token* type_tok = advance(&p);  // 跳过 'type'
            const Token* alias_name = expect(&p, TOK_IDENT, error, "type 后应为别名");
            if (!alias_name) return NULL;
            if (!expect(&p, TOK_ASSIGN, error, "缺少 '='")) return NULL;
            TypeInfo rti = parse_type_expr(&p, error);
            if (rti.kind == TYPE_ERROR) return NULL;
            if (!expect(&p, TOK_SEMICOLON, error, "缺少 ';'")) return NULL;
            aliases[alias_count++] = ast_make_type_alias(a,
                arena_strdup_impl(a, alias_name->start, alias_name->length),
                rti.kind, rti.struct_name, tok_loc(type_tok));
        } else if (peek(&p)->kind == TOK_ENUM) {
            advance(&p);
            const Token* name = expect(&p, TOK_IDENT, error, "enum 后应为枚举名");
            if (!name) return NULL;
            if (!expect(&p, TOK_LBRACE, error, "缺少 '{'")) return NULL;
            const char* variants[64];
            TypeKind payload_types[64];
            const char* payload_snames[64];
            int vcount = 0;
            while (peek(&p)->kind != TOK_RBRACE && !error->message) {
                if (vcount >= 64) { error->message = "枚举变体过多(最多64)"; error->filename = p.filename; error->line = peek(&p)->line; error->col = peek(&p)->col; return NULL; }
                const Token* vname = expect(&p, TOK_IDENT, error, "变体名");
                if (!vname) return NULL;
                variants[vcount] = arena_strdup_impl(p.arena, vname->start, vname->length);
                payload_types[vcount] = TYPE_VOID;
                payload_snames[vcount] = NULL;
                // 可选 payload: Variant(Type)
                if (peek(&p)->kind == TOK_LPAREN) {
                    advance(&p);
                    TypeInfo pti = parse_type_expr(&p, error);
                    if (pti.kind == TYPE_ERROR) return NULL;
                    payload_types[vcount] = pti.kind;
                    payload_snames[vcount] = pti.struct_name;
                    if (!expect(&p, TOK_RPAREN, error, "缺少 ')'")) return NULL;
                }
                vcount++;
                if (peek(&p)->kind == TOK_COMMA) advance(&p); else break;
            }
            if (!expect(&p, TOK_RBRACE, error, "缺少 '}'")) return NULL;
            const char** v_arr = arena_alloc_impl(p.arena, vcount * sizeof(const char*));
            memcpy(v_arr, variants, vcount * sizeof(const char*));
            TypeKind* pt_arr = arena_alloc_impl(p.arena, vcount * sizeof(TypeKind));
            memcpy(pt_arr, payload_types, vcount * sizeof(TypeKind));
            const char** ps_arr = arena_alloc_impl(p.arena, vcount * sizeof(const char*));
            memcpy(ps_arr, payload_snames, vcount * sizeof(const char*));
            AstNode* enum_decl = ast_make_enum_decl(p.arena,
                arena_strdup_impl(p.arena, name->start, name->length),
                v_arr, pt_arr, ps_arr, vcount, tok_loc(name));
            if (enum_count >= 64) { error->message = "枚举过多 (最多64)"; error->filename = p.filename; error->line = peek(&p)->line; error->col = peek(&p)->col; return NULL; }
            enums[enum_count++] = enum_decl;
        } else if (peek(&p)->kind == TOK_EXTEND) {
            const Token* i_tok = advance(&p);
            const Token* first = expect(&p, TOK_IDENT, error, "extend 后应为结构体名");
            if (!first) return NULL;
            const char* trait_name = NULL;
            const char* struct_name;
            // extend Trait Struct { ... }（trait 实现：两个标识符）
            if (peek(&p)->kind == TOK_IDENT) {
                trait_name = arena_strdup_impl(p.arena, first->start, first->length);
                struct_name = arena_strdup_impl(p.arena, peek(&p)->start, peek(&p)->length);
                advance(&p);
            } else {
                struct_name = arena_strdup_impl(p.arena, first->start, first->length);
            }
            if (!expect(&p, TOK_LBRACE, error, "缺少 '{'")) return NULL;
            AstNode* methods[64]; int mcount = 0;
            while (peek(&p)->kind != TOK_RBRACE && !error->message) {
                if (mcount >= 64) { error->message = "方法过多 (最多64)"; error->filename = p.filename; error->line = peek(&p)->line; error->col = peek(&p)->col; return NULL; }
                if (peek(&p)->kind != TOK_FN) { error->message = "extend 块内只允许 fn"; error->filename = p.filename; error->line = peek(&p)->line; error->col = peek(&p)->col; return NULL; }
                AstNode* m = parse_function(&p, false, error);
                if (!m) return NULL;
                // trait 实现: 方法名 mangled 为 TraitName$methodName
                if (trait_name) {
                    char* mn = arena_alloc_impl(p.arena, strlen(trait_name) + strlen(m->as.function.name) + 4);
                    sprintf(mn, "%s$%s", trait_name, m->as.function.name);
                    m->as.function.name = mn;
                }
                methods[mcount++] = m;
            }
            if (!expect(&p, TOK_RBRACE, error, "缺少 '}'")) return NULL;
            AstNode** m_arr = arena_alloc_impl(p.arena, mcount * sizeof(AstNode*));
            memcpy(m_arr, methods, mcount * sizeof(AstNode*));
            if (impl_count >= 64) { error->message = "extend 块过多 (最多64)"; error->filename = p.filename; error->line = peek(&p)->line; error->col = peek(&p)->col; return NULL; }
            impls[impl_count++] = ast_make_impl_block(p.arena, struct_name, m_arr, mcount, tok_loc(i_tok));
        } else if (peek(&p)->kind == TOK_MOD) {
            advance(&p);
            const Token* mn = expect(&p, TOK_IDENT, error, "mod 后应为模块名");
            if (!mn) return NULL;
            if (!expect(&p, TOK_SEMICOLON, error, "缺少 ';'")) return NULL;
            const char* mod_name = arena_strdup_impl(p.arena, mn->start, mn->length);
            AstNode* sub = load_module(a, filename, mod_name, error);
            if (!sub) return NULL;
            // 合并子模块项到当前文件（以 mod_name:: 为前缀）
            for (size_t i = 0; i < sub->as.program.fn_count; i++) {
                AstNode* fn = sub->as.program.functions[i];
                if (fn->as.function.is_pub) {
                    char* mangled = arena_alloc_impl(p.arena, strlen(mod_name) + strlen(fn->as.function.name) + 4);
                    sprintf(mangled, "%s::%s", mod_name, fn->as.function.name);
                    fn->as.function.name = mangled;
                    if (fn_count >= 256) { error->message = "函数过多"; error->filename = p.filename; return NULL; }
                    functions[fn_count++] = fn;
                }
            }
            for (size_t i = 0; i < sub->as.program.struct_count; i++) {
                if (struct_count >= 64) break;
                structs[struct_count++] = sub->as.program.structs[i];
            }
            for (size_t i = 0; i < sub->as.program.enum_count; i++) {
                if (enum_count >= 64) break;
                enums[enum_count++] = sub->as.program.enums[i];
            }
        } else if (peek(&p)->kind == TOK_USE) {
            /* TODO: use 语句待实现符号导入 */ ;
            advance(&p);
            while (peek(&p)->kind != TOK_SEMICOLON && peek(&p)->kind != TOK_EOF) advance(&p);
            if (peek(&p)->kind == TOK_SEMICOLON) advance(&p);
        } else if (peek(&p)->kind == TOK_FN) {
            if (fn_count >= 256) { error->message = "函数过多 (最多256)"; error->filename = p.filename; error->line = peek(&p)->line; error->col = peek(&p)->col; return NULL; }
            functions[fn_count++] = parse_function(&p, is_pub, error);
        } else {
            error->message = "顶层只允许 fn、struct、type、enum、extend、mod 或 use";
            error->filename = p.filename;
            error->line = peek(&p)->line;
            error->col = peek(&p)->col;
            return NULL;
        }
    }
    if (error->message) return NULL;
    AstNode** fn_arr = arena_alloc_impl(a, fn_count * sizeof(AstNode*));
    memcpy(fn_arr, functions, fn_count * sizeof(AstNode*));
    AstNode** st_arr = arena_alloc_impl(a, struct_count * sizeof(AstNode*));
    memcpy(st_arr, structs, struct_count * sizeof(AstNode*));
    AstNode** al_arr = arena_alloc_impl(a, alias_count * sizeof(AstNode*));
    memcpy(al_arr, aliases, alias_count * sizeof(AstNode*));
    AstNode** en_arr = arena_alloc_impl(a, enum_count * sizeof(AstNode*));
    memcpy(en_arr, enums, enum_count * sizeof(AstNode*));
    AstNode** im_arr = arena_alloc_impl(a, impl_count * sizeof(AstNode*));
    memcpy(im_arr, impls, impl_count * sizeof(AstNode*));
    return ast_make_program(a, fn_arr, fn_count, st_arr, struct_count,
                            al_arr, alias_count, en_arr, enum_count,
                            im_arr, impl_count, loc_at(0, 0));
}