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

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

spl_tok_t *peek(spl_comp_t *ctx) { return &vec_at(ctx->toks, ctx->tok_idx); }

spl_tok_t *advance(spl_comp_t *ctx) {
    spl_tok_t *t = &vec_at(ctx->toks, ctx->tok_idx);
    if (t->type != TOK_EOF)
        ctx->tok_idx++;
    return t;
}


/* ============================================================
 * Parse function definition
 *   fn name(params) ret-type { body }
 *   or fn name(params) ret-type;  (forward decl, not used for stage1)
 * ============================================================ */

/* Shared helper: register a function, declare params, parse body, end function.
 * Used by both top-level fn decl and methods inside type bodies. */
static int parse_fn_body(spl_comp_t *ctx, const char *fn_name, spl_type_info_t *ret_type,
                         int nparams, char pnames[][256], spl_type_info_t *ptypes[], int is_pub) {
    int fi = spl_declare_func(ctx, fn_name, ret_type, nparams, 0, is_pub);
    ctx->current_func_idx = fi;
    ctx->current_ret_type = ret_type;
    ctx->current_local_bytes = 0;

    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);
        spl_emit(ctx, SPL_ALLOC, SPL_VOID, 0);

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

        spl_patch(ctx, alloc_addr, ctx->current_local_bytes / (int)sizeof(spl_val_t) - nparams);
        expect(ctx, TOK_R_BRACE);
    }

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

    spl_emit(ctx, SPL_RET, SPL_VOID, 0);
    spl_prog_end_func(&ctx->prog, fi);
    ctx->current_func_idx = -1;
    ctx->current_ret_type = NULL;
    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];
    usize fnl = fname_tok->len < 255 ? fname_tok->len : 255;
    memcpy(fn_name, fname_tok->lexeme, fnl);
    fn_name[fnl] = '\0';

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

    /* Parse parameters — collect names and types */
    int nparams = 0;
    enum { MAX_PARAMS = 64 };
    char pnames[MAX_PARAMS][256];
    spl_type_info_t *ptypes[MAX_PARAMS];
    skip_nl(ctx);
    if (peek(ctx)->type != TOK_R_PAREN) {
        while (1) {
            spl_tok_t *pname = advance(ctx);
            usize pnl = pname->len < 255 ? pname->len : 255;
            memcpy(pnames[nparams], pname->lexeme, pnl);
            pnames[nparams][pnl] = '\0';
            skip_nl(ctx);
            if (peek(ctx)->type == TOK_COLON) {
                advance(ctx); /* : */
                skip_nl(ctx);
                ptypes[nparams] = spl_parse_type(ctx);
            } else {
                ptypes[nparams] = spl_type_basic(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);
    skip_nl(ctx);

    /* Return type (default: void) */
    spl_type_info_t *ret_type = spl_type_basic(SPL_VOID);
    if (peek(ctx)->type != TOK_SEMICOLON && peek(ctx)->type != TOK_L_BRACE) {
        ret_type = spl_parse_type(ctx);
        if (!ret_type)
            ret_type = spl_type_basic(SPL_VOID);
        skip_nl(ctx);
    }

    /* Extern function: register as native, no body */
    if (is_extern) {
        spl_declare_func(ctx, fn_name, ret_type, nparams, 1, is_pub);
        /* Add to prog->natives for NCALL dispatch */
        int found = -1;
        vec_for(ctx->prog.natives, ni) {
            if (strcmp(vec_at(ctx->prog.natives, ni).name, fn_name) == 0) {
                found = (int)ni;
                break;
            }
        }
        if (found < 0) {
            spl_native_t nat;
            memset(&nat, 0, sizeof(nat));
            nat.name = strdup(fn_name);
            nat.idx_of_strtab = 0;
            nat.impl_fn = NULL; /* resolved by VM at runtime */
            vec_push(ctx->prog.natives, nat);
        }
        if (peek(ctx)->type == TOK_SEMICOLON)
            advance(ctx);
        return;
    }

    /* Check for forward declaration (just semicolon, skip) */
    if (peek(ctx)->type == TOK_SEMICOLON) {
        advance(ctx);
        return;
    }

    skip_nl(ctx);

    /* Use shared helper for function body parsing */
    parse_fn_body(ctx, fn_name, ret_type, nparams, pnames, ptypes, is_pub);
}

/* ============================================================
 * Parse struct/union body (shared for struct and union containers)
 *
 * Body supports:
 *   var name: type;       — field declarations
 *   name: type,           — field declarations (old-style)
 *   type Name = ...;      — nested type declarations
 *   fn name(...) type { } — methods
 * ============================================================ */

static void parse_struct_body(spl_comp_t *ctx, spl_type_info_t *st) {
    if (peek(ctx)->type != TOK_L_BRACE)
        return;
    advance(ctx); /* { */

    /* === Pass 1: Parse all nested type declarations first ===
     * This allows field declarations to reference types defined later in the body. */
    {
        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);
            }
        }
        /* Reset to start of body for second pass */
        ctx->tok_idx = saved;
    }

    /* === Pass 2: Parse fields 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) {
                /* Nested type — already parsed in pass 1, skip by re-parsing */
                parse_type_decl(ctx);
            } else if (tt == KW_VAR && depth == 1) {
                /* var name: type; */
                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);
                    spl_type_info_t *ftype = spl_parse_type(ctx);
                    char fname[256];
                    usize fnl = ftok->len < 255 ? ftok->len : 255;
                    memcpy(fname, ftok->lexeme, fnl);
                    fname[fnl] = '\0';
                    spl_type_add_field(st, 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) {
                /* Old-style field: name: type, */
                spl_tok_t *ftok = advance(ctx);
                skip_nl(ctx);
                if (peek(ctx)->type == TOK_COLON) {
                    advance(ctx); /* : */
                    skip_nl(ctx);
                    spl_type_info_t *ftype = spl_parse_type(ctx);
                    char fname[256];
                    usize fnl = ftok->len < 255 ? ftok->len : 255;
                    memcpy(fname, ftok->lexeme, fnl);
                    fname[fnl] = '\0';
                    spl_type_add_field(st, 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) {
                /* Method — parse properly using parse_fn_body */
                advance(ctx); /* fn */
                skip_nl(ctx);
                spl_tok_t *mname_tok = advance(ctx);
                char mname[256];
                usize mnl = mname_tok->len < 255 ? mname_tok->len : 255;
                memcpy(mname, mname_tok->lexeme, mnl);
                mname[mnl] = '\0';

                /* Build qualified name: TypeName.method_name */
                char qualified[512];
                snprintf(qualified, sizeof(qualified), "%s.%s", st->name ? st->name : "anon",
                         mname);

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

                /* Parse parameters */
                int nparams = 0;
                enum { MAX_PARAMS = 64 };
                char pnames[MAX_PARAMS][256];
                spl_type_info_t *ptypes[MAX_PARAMS];
                skip_nl(ctx);
                if (peek(ctx)->type != TOK_R_PAREN) {
                    while (1) {
                        spl_tok_t *pname = advance(ctx);
                        usize pnl = pname->len < 255 ? pname->len : 255;
                        memcpy(pnames[nparams], pname->lexeme, pnl);
                        pnames[nparams][pnl] = '\0';
                        skip_nl(ctx);
                        if (peek(ctx)->type == TOK_COLON) {
                            advance(ctx); /* : */
                            skip_nl(ctx);
                            ptypes[nparams] = spl_parse_type(ctx);
                        } else {
                            ptypes[nparams] = spl_type_basic(SPL_I32);
                        }
                        nparams++;
                        skip_nl(ctx);
                        if (peek(ctx)->type == TOK_COMMA) {
                            advance(ctx);
                            skip_nl(ctx);
                            continue;
                        }
                        break;
                    }
                }
                expect(ctx, TOK_R_PAREN);
                skip_nl(ctx);

                /* Return type (default: void) */
                spl_type_info_t *ret_type = spl_type_basic(SPL_VOID);
                if (peek(ctx)->type != TOK_SEMICOLON && peek(ctx)->type != TOK_L_BRACE) {
                    ret_type = spl_parse_type(ctx);
                    if (!ret_type)
                        ret_type = spl_type_basic(SPL_VOID);
                    skip_nl(ctx);
                }

                /* Set current_type_name for short-name resolution inside method body */
                const char *saved_type_name = ctx->current_type_name;
                ctx->current_type_name = st->name;

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

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

                ctx->current_type_name = saved_type_name;

                spl_type_add_method(st, mname, fi);
            } else {
                advance(ctx);
            }
        }
    }

    spl_type_compute_layout(st);
}

/* ============================================================
 * Parse enum body
 *
 * Body supports:
 *   Name                  — simple variant
 *   Name: Type            — variant with data
 *   var name: type;       — field-style variant
 *   type Name = ...;      — nested type declarations
 *   fn name(...) type { } — methods (skipped)
 * ============================================================ */

static void parse_enum_body(spl_comp_t *ctx, spl_type_info_t *et) {
    if (peek(ctx)->type != TOK_L_BRACE)
        return;
    advance(ctx); /* { */

    /* === 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;
    }

    /* === Pass 2: Parse 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) {
                /* Already parsed in pass 1, re-parse to skip */
                parse_type_decl(ctx);
            } else if (tt == TOK_IDENT && depth == 1) {
                /* Variant: Name or Name: Type */
                spl_tok_t *vtok = advance(ctx);
                skip_nl(ctx);
                if (peek(ctx)->type == TOK_COLON) {
                    advance(ctx); /* : */
                    skip_nl(ctx);
                    spl_type_info_t *dtype = spl_parse_type(ctx);
                    char vname[256];
                    usize vnl = vtok->len < 255 ? vtok->len : 255;
                    memcpy(vname, vtok->lexeme, vnl);
                    vname[vnl] = '\0';
                    spl_type_add_variant(et, vname, dtype);
                } else {
                    char vname[256];
                    usize vnl = vtok->len < 255 ? vtok->len : 255;
                    memcpy(vname, vtok->lexeme, vnl);
                    vname[vnl] = '\0';
                    spl_type_add_variant(et, vname, NULL);
                }
                skip_nl(ctx);
                if (peek(ctx)->type == TOK_SEMICOLON || peek(ctx)->type == TOK_COMMA)
                    advance(ctx);
            } else if (tt == KW_FN && depth == 1) {
                /* Method — parse properly using parse_fn_body */
                advance(ctx); /* fn */
                skip_nl(ctx);
                spl_tok_t *mname_tok = advance(ctx);
                char mname[256];
                usize mnl = mname_tok->len < 255 ? mname_tok->len : 255;
                memcpy(mname, mname_tok->lexeme, mnl);
                mname[mnl] = '\0';

                /* Build qualified name: TypeName.method_name */
                char qualified[512];
                snprintf(qualified, sizeof(qualified), "%s.%s", et->name ? et->name : "anon",
                         mname);

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

                /* Parse parameters */
                int nparams = 0;
                enum { MAX_PARAMS = 64 };
                char pnames[MAX_PARAMS][256];
                spl_type_info_t *ptypes[MAX_PARAMS];
                skip_nl(ctx);
                if (peek(ctx)->type != TOK_R_PAREN) {
                    while (1) {
                        spl_tok_t *pname = advance(ctx);
                        usize pnl = pname->len < 255 ? pname->len : 255;
                        memcpy(pnames[nparams], pname->lexeme, pnl);
                        pnames[nparams][pnl] = '\0';
                        skip_nl(ctx);
                        if (peek(ctx)->type == TOK_COLON) {
                            advance(ctx); /* : */
                            skip_nl(ctx);
                            ptypes[nparams] = spl_parse_type(ctx);
                        } else {
                            ptypes[nparams] = spl_type_basic(SPL_I32);
                        }
                        nparams++;
                        skip_nl(ctx);
                        if (peek(ctx)->type == TOK_COMMA) {
                            advance(ctx);
                            skip_nl(ctx);
                            continue;
                        }
                        break;
                    }
                }
                expect(ctx, TOK_R_PAREN);
                skip_nl(ctx);

                /* Return type (default: void) */
                spl_type_info_t *ret_type = spl_type_basic(SPL_VOID);
                if (peek(ctx)->type != TOK_SEMICOLON && peek(ctx)->type != TOK_L_BRACE) {
                    ret_type = spl_parse_type(ctx);
                    if (!ret_type)
                        ret_type = spl_type_basic(SPL_VOID);
                    skip_nl(ctx);
                }

                /* Set current_type_name for short-name resolution inside method body */
                const char *saved_type_name = ctx->current_type_name;
                ctx->current_type_name = et->name;

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

                ctx->current_type_name = saved_type_name;

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

                spl_type_add_method(et, mname, fi);
            } else {
                advance(ctx);
            }
        }
    }

    spl_type_compute_layout(et);
}

/* ============================================================
 * Parse type declaration
 *   type Name = struct { ... };
 *   type Name = union { ... };
 *   type Name = enum { ... };
 *   type Name = ExistingType;
 * ============================================================ */

void parse_type_decl(spl_comp_t *ctx) {
    advance(ctx); /* type */
    spl_tok_t *name_tok = advance(ctx);
    char tname[256];
    usize tnl = name_tok->len < 255 ? name_tok->len : 255;
    memcpy(tname, name_tok->lexeme, tnl);
    tname[tnl] = '\0';

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

    if (peek(ctx)->type == KW_STRUCT) {
        advance(ctx);
        spl_type_info_t *st = spl_type_struct(tname);
        /* Register type early to allow self-referential fields */
        map_put(ctx->type_defs, strdup(tname), st);
        skip_nl(ctx);
        parse_struct_body(ctx, st);
    } else if (peek(ctx)->type == KW_UNION) {
        advance(ctx);
        spl_type_info_t *ut = spl_type_union(tname);
        map_put(ctx->type_defs, strdup(tname), ut);
        skip_nl(ctx);
        parse_struct_body(ctx, ut);
    } else if (peek(ctx)->type == KW_ENUM) {
        advance(ctx);
        spl_type_info_t *et = spl_type_enum(tname);
        /* Register type early to allow self-referential variants */
        map_put(ctx->type_defs, strdup(tname), et);
        skip_nl(ctx);
        parse_enum_body(ctx, et);
    } else if (peek(ctx)->type == TOK_IDENT ||
               (peek(ctx)->type >= KW_AS && peek(ctx)->type <= KW_ANY)) {
        spl_type_info_t *base = spl_parse_type(ctx);
        if (base) {
            spl_type_info_t *alias = spl_type_clone(base);
            alias->name = strdup(tname);
            alias->kind = TYPE_NAME;
            map_put(ctx->type_defs, strdup(tname), alias);
        }
    }

    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_SHARP:
            /* #[extern("vm")] fn ... */
            advance(ctx); /* # */
            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;
            /* Try to parse as a statement */
            spl_parse_stmt(ctx);
            /* Safety: prevent infinite loop on unrecognized tokens */
            if (ctx->tok_idx == prev)
                advance(ctx);
            break;
        }
        }
    }
}