/* spl_parser.c — Top-level parser: function declarations, type declarations, etc. */

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

/* ============================================================
 * Parse function definition
 * ============================================================ */

enum { MAX_PARAMS = 64 };

static int parse_params_decl(spl_comp_t *ctx, char pnames[][256], int ptypes[]) {
    int nparams = 0;
    skip_nl(ctx);
    if (peek(ctx)->type != TOK_R_PAREN) {
        while (1) {
            spl_tok_t *pname = advance(ctx);
            spl_tok_copy_name(pname, pnames[nparams], 256);
            skip_nl(ctx);
            if (peek(ctx)->type == TOK_COLON) {
                advance(ctx);
                skip_nl(ctx);
                ptypes[nparams] = spl_type_parse(&ctx->tctx, ctx);
            } else {
                ptypes[nparams] = spl_type_basic(&ctx->tctx, SPL_I32);
            }
            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);
            }
            break;
        }
    }
    expect(ctx, TOK_R_PAREN);
    return nparams;
}

static int parse_fn_body(spl_comp_t *ctx, const char *fn_name, int ret_type_idx, int nparams,
                         char pnames[][256], int ptypes[], int is_pub) {
    int fi = spl_declare_func(ctx, fn_name, ret_type_idx, nparams, 0, is_pub);
    ctx->current_func_idx = fi;
    ctx->current_ret_type_idx = ret_type_idx;
    fa_init(&ctx->emit.frame);

    spl_push_scope(ctx);

    for (int i = 0; i < nparams; i++) {
        spl_declare_var(ctx, pnames[i], ptypes[i], 0);
    }

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

        spl_val_t alloc_addr = vec_size(ctx->prog.insns);
        emit_alloc(&ctx->emit, 0);

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

        int total_phys_slots = 0;
        for (int i = 0; i < nparams; i++) {
            usize psz = spl_type_size(&ctx->tctx, ptypes[i]);
            total_phys_slots += (int)((psz + sizeof(spl_val_t) - 1) / sizeof(spl_val_t));
        }
        if (spl_type_needs_multi_slot(&ctx->tctx, ret_type_idx)) {
            usize min_bytes = spl_type_slot_count(&ctx->tctx, ret_type_idx) * sizeof(spl_val_t);
            if ((usize)ctx->emit.frame.peak_bytes < min_bytes)
                ctx->emit.frame.peak_bytes = (int)min_bytes;
        }
        emit_patch(&ctx->emit, alloc_addr,
                   ctx->emit.frame.peak_bytes / (int)sizeof(spl_val_t) - total_phys_slots);
        expect(ctx, TOK_R_BRACE);
    }

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

    {
        spl_func_info_t *f = &vec_at(ctx->funcs, fi);
        f->param_type_indices = calloc(nparams, sizeof(int));
        f->param_names = calloc(nparams, sizeof(char *));
        for (int i = 0; i < nparams; i++) {
            f->param_type_indices[i] = ptypes[i];
            f->param_names[i] = strdup(pnames[i]);
        }
    }

    emit_return(&ctx->emit, &ctx->tctx, ctx->current_ret_type_idx);
    spl_prog_end_func(&ctx->prog, fi);
    ctx->current_func_idx = -1;
    ctx->current_ret_type_idx = -1;
    return fi;
}

static void parse_fn_decl(spl_comp_t *ctx, int is_extern, int is_pub) {
    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);

    char pnames[MAX_PARAMS][256];
    int ptypes[MAX_PARAMS];
    int nparams = parse_params_decl(ctx, pnames, ptypes);
    skip_nl(ctx);

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

    int fi;

    if (is_extern) {
        fi = spl_declare_func(ctx, fn_name, ret_type_idx, nparams, 1, is_pub);
        spl_ensure_native(ctx, fn_name);
        if (ctx->tctx.current_type_idx == 0)
            spl_type_add_method(&ctx->tctx, 0, fn_name, fi);
        if (peek(ctx)->type == TOK_SEMICOLON)
            advance(ctx);
        return;
    }

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

    skip_nl(ctx);

    fi = parse_fn_body(ctx, fn_name, ret_type_idx, nparams, pnames, ptypes, is_pub);
    if (ctx->tctx.current_type_idx == 0)
        spl_type_add_method(&ctx->tctx, 0, fn_name, fi);
}

/* ============================================================
 * Parse method declaration inside a type body
 * ============================================================ */

static void parse_method_decl(spl_comp_t *ctx, int container_type_idx) {
    advance(ctx); /* fn */
    skip_nl(ctx);
    spl_tok_t *mname_tok = advance(ctx);
    char mname[256];
    spl_tok_copy_name(mname_tok, mname, sizeof(mname));
    const char *cname = spl_type_name(&ctx->tctx, container_type_idx);
    char qualified[512];
    snprintf(qualified, sizeof(qualified), "%s.%s", cname ? cname : "anon", mname);

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

    char pnames[MAX_PARAMS][256];
    int ptypes[MAX_PARAMS];
    int nparams = parse_params_decl(ctx, pnames, ptypes);
    skip_nl(ctx);

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

    int fi = parse_fn_body(ctx, qualified, ret_type_idx, nparams, pnames, ptypes, 0);

    spl_type_add_method(&ctx->tctx, container_type_idx, mname, fi);
}

/* ============================================================
 * Parse type container body (shared for struct, union, enum)
 * ============================================================ */

static void skip_type_decl(spl_comp_t *ctx) {
    advance(ctx); /* type */
    advance(ctx); /* name */
    skip_nl(ctx);
    expect(ctx, TOK_ASSIGN);
    skip_nl(ctx);

    spl_tok_type_t tt = peek(ctx)->type;
    if (tt == KW_STRUCT || tt == KW_UNION || tt == KW_ENUM) {
        advance(ctx);
        skip_nl(ctx);
    }

    if (peek(ctx)->type == TOK_L_BRACE) {
        int bd = 1;
        advance(ctx);
        while (bd > 0 && ctx->tok_idx < vec_size(ctx->toks)) {
            if (peek(ctx)->type == TOK_L_BRACE)
                bd++;
            else if (peek(ctx)->type == TOK_R_BRACE)
                bd--;
            advance(ctx);
        }
    } else {
        while (peek(ctx)->type != TOK_SEMICOLON && peek(ctx)->type != TOK_EOF)
            advance(ctx);
    }

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

static void parse_type_body(spl_comp_t *ctx, int container_type_idx, int is_enum) {
    if (peek(ctx)->type != TOK_L_BRACE)
        return;
    advance(ctx); /* { */

    int saved_current = ctx->tctx.current_type_idx;
    ctx->tctx.current_type_idx = container_type_idx;

    /* === Pass 1: Parse all nested type declarations first === */
    {
        usize saved = ctx->tok_idx;
        int depth = 1;
        while (depth > 0 && ctx->tok_idx < vec_size(ctx->toks)) {
            spl_tok_type_t tt = peek(ctx)->type;
            if (tt == TOK_L_BRACE) {
                depth++;
                advance(ctx);
            } else if (tt == TOK_R_BRACE) {
                depth--;
                if (depth == 0)
                    break;
                advance(ctx);
            } else if (tt == KW_TYPE && depth == 1) {
                parse_type_decl(ctx);
            } else {
                advance(ctx);
            }
        }
        ctx->tok_idx = saved;
    }

    /* === Pre-register method names for forward references === */
    {
        const char *cname = spl_type_name(&ctx->tctx, container_type_idx);
        if (cname) {
            usize saved = ctx->tok_idx;
            int depth = 1;
            while (depth > 0 && ctx->tok_idx < vec_size(ctx->toks)) {
                spl_tok_type_t tt = peek(ctx)->type;
                if (tt == TOK_L_BRACE) {
                    depth++;
                    advance(ctx);
                } else if (tt == TOK_R_BRACE) {
                    depth--;
                    if (depth == 0)
                        break;
                    advance(ctx);
                } else if (tt == KW_FN && depth == 1) {
                    advance(ctx); /* fn */
                    skip_nl(ctx);
                    spl_tok_t *name_tok = advance(ctx);
                    char mname[256];
                    spl_tok_copy_name(name_tok, mname, sizeof(mname));
                    char qualified[512];
                    snprintf(qualified, sizeof(qualified), "%s.%s", cname, mname);
                    int fi = spl_declare_func(ctx, qualified, -1, 0, 0, 0);
                    spl_type_add_method(&ctx->tctx, container_type_idx, mname, fi);
                } else {
                    advance(ctx);
                }
            }
            ctx->tok_idx = saved;
        }
    }

    /* === Pass 2: Parse fields/variants and methods === */
    {
        int depth = 1;
        while (depth > 0 && ctx->tok_idx < vec_size(ctx->toks)) {
            skip_nl(ctx);
            spl_tok_type_t tt = peek(ctx)->type;
            if (tt == TOK_L_BRACE) {
                depth++;
                advance(ctx);
            } else if (tt == TOK_R_BRACE) {
                depth--;
                if (depth == 0) {
                    advance(ctx);
                    break;
                }
                advance(ctx);
            } else if (tt == KW_TYPE && depth == 1) {
                skip_type_decl(ctx);
            } else if (tt == KW_VAR && depth == 1 && !is_enum) {
                advance(ctx); /* var */
                skip_nl(ctx);
                spl_tok_t *ftok = advance(ctx);
                skip_nl(ctx);
                if (peek(ctx)->type == TOK_COLON) {
                    advance(ctx); /* : */
                    skip_nl(ctx);
                    int ftype = spl_type_parse(&ctx->tctx, ctx);
                    char fname[256];
                    spl_tok_copy_name(ftok, fname, sizeof(fname));
                    if (ftype >= 0)
                        spl_type_add_field(&ctx->tctx, container_type_idx, fname, ftype);
                }
                skip_nl(ctx);
                if (peek(ctx)->type == TOK_SEMICOLON || peek(ctx)->type == TOK_COMMA)
                    advance(ctx);
            } else if (tt == TOK_IDENT && depth == 1) {
                if (is_enum) {
                    spl_tok_t *vtok = advance(ctx);
                    skip_nl(ctx);
                    if (peek(ctx)->type == TOK_COLON) {
                        advance(ctx); /* : */
                        skip_nl(ctx);
                        int dtype = spl_type_parse(&ctx->tctx, ctx);
                        char vname[256];
                        spl_tok_copy_name(vtok, vname, sizeof(vname));
                        spl_type_add_variant(&ctx->tctx, container_type_idx, vname,
                                             dtype >= 0 ? dtype : -1);
                    } else {
                        char vname[256];
                        spl_tok_copy_name(vtok, vname, sizeof(vname));
                        spl_type_add_variant(&ctx->tctx, container_type_idx, vname, -1);
                    }
                } else {
                    spl_tok_t *ftok = advance(ctx);
                    skip_nl(ctx);
                    if (peek(ctx)->type == TOK_COLON) {
                        advance(ctx); /* : */
                        skip_nl(ctx);
                        int ftype = spl_type_parse(&ctx->tctx, ctx);
                        char fname[256];
                        spl_tok_copy_name(ftok, fname, sizeof(fname));
                        if (ftype >= 0)
                            spl_type_add_field(&ctx->tctx, container_type_idx, fname, ftype);
                    }
                }
                skip_nl(ctx);
                if (peek(ctx)->type == TOK_SEMICOLON || peek(ctx)->type == TOK_COMMA)
                    advance(ctx);
            } else if (tt == KW_FN && depth == 1) {
                spl_type_compute_layout(&ctx->tctx, container_type_idx);
                parse_method_decl(ctx, container_type_idx);
            } else {
                advance(ctx);
            }
        }
    }

    ctx->tctx.current_type_idx = saved_current;
    spl_type_compute_layout(&ctx->tctx, container_type_idx);
}

/* ============================================================
 * Parse type declaration
 * ============================================================ */

void parse_type_decl(spl_comp_t *ctx) {
    advance(ctx); /* type */
    spl_tok_t *name_tok = advance(ctx);
    char tname[256];
    spl_tok_copy_name(name_tok, tname, sizeof(tname));

    skip_nl(ctx);
    expect(ctx, TOK_ASSIGN);
    skip_nl(ctx);

    int parent_type_idx = ctx->tctx.current_type_idx;

    if (peek(ctx)->type == KW_STRUCT) {
        advance(ctx);
        int ti = spl_type_struct(&ctx->tctx, tname);
        if (parent_type_idx >= 0)
            spl_type_add_nested(&ctx->tctx, parent_type_idx, tname, ti);
        skip_nl(ctx);
        parse_type_body(ctx, ti, 0);
    } else if (peek(ctx)->type == KW_UNION) {
        advance(ctx);
        int ti = spl_type_union(&ctx->tctx, tname);
        if (parent_type_idx >= 0)
            spl_type_add_nested(&ctx->tctx, parent_type_idx, tname, ti);
        skip_nl(ctx);
        parse_type_body(ctx, ti, 1);
    } else if (peek(ctx)->type == KW_ENUM) {
        advance(ctx);
        int ti = spl_type_enum(&ctx->tctx, tname);
        if (parent_type_idx >= 0)
            spl_type_add_nested(&ctx->tctx, parent_type_idx, tname, ti);
        skip_nl(ctx);
        parse_type_body(ctx, ti, 1);
    } else if (peek(ctx)->type == TOK_IDENT ||
               (peek(ctx)->type >= KW_AS && peek(ctx)->type <= KW_ANY)) {
        int base = spl_type_parse(&ctx->tctx, ctx);
        if (base >= 0) {
            spl_type_alias(&ctx->tctx, tname, base);
        }
    }

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

/* ============================================================
 * Parse top-level program
 * ============================================================ */

void spl_parse_prog(spl_comp_t *ctx) {
    while (peek(ctx)->type != TOK_EOF) {
        skip_nl(ctx);
        if (peek(ctx)->type == TOK_EOF)
            break;

        switch (peek(ctx)->type) {
        case KW_FN:
            parse_fn_decl(ctx, 0, 0);
            break;
        case KW_TYPE:
            parse_type_decl(ctx);
            break;
        case KW_PUB: {
            advance(ctx); /* skip pub */
            skip_nl(ctx);
            if (peek(ctx)->type == KW_FN)
                parse_fn_decl(ctx, 0, 1);
            else if (peek(ctx)->type == KW_TYPE)
                parse_type_decl(ctx);
            break;
        }
        case TOK_AT:
        case TOK_SHARP: {
            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 (e.g. "vm") */
                        skip_nl(ctx);
                        if (peek(ctx)->type == TOK_R_PAREN)
                            advance(ctx);
                    }
                }
            } else {
                /* #[extern("vm")] fn ... (legacy) */
                if (peek(ctx)->type == TOK_L_BRACKET) {
                    advance(ctx); /* [ */
                    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)
                parse_fn_decl(ctx, 1, 0);
            break;
        }
        default: {
            int prev = ctx->tok_idx;
            spl_parse_stmt(ctx);
            if (ctx->tok_idx == prev)
                advance(ctx);
            break;
        }
        }
    }
}