#include "parse.h"
#include "ast.h"
#include "collection.h"
#include "syms.h"
#include <assert.h>
#include <stddef.h>
#include <stdio.h>
#include <stdlib.h>

typedef struct {
    char ch;
    bool is_term;
    unsigned short val;
} KwNodeOpt;

typedef struct {
    ARRAY_FIELD(KwNodeOpt, opts);
    size_t idx;
} KwNode;

typedef struct {
    TokStream* tokstream;
    Arena arena;
    ARRAY_FIELD(KwNode, nodes);
    KwNode* root;
} KwMatcher;

static KwNodeOpt* kwmatcher_create_path(
    KwMatcher* m, KwNode* node, const char* kw)
{
    assert(kw[0]);
    KwNodeOpt* opt = nullptr;
    for (size_t i = 0; i < node->opts_size; ++i) {
        if (node->opts[i].ch == kw[0]) {
            opt = &node->opts[i];
        }
    }
    if (!opt) {
        size_t new_idx = m->nodes_size;

        KwNode new_node = {
            ARRAY_FIELD_INIT(opts),
            .idx = new_idx,
        };
        arena_array_push(&m->arena,
            (void**)&m->nodes,
            &m->nodes_cap,
            &m->nodes_size,
            &new_node,
            sizeof(new_node));
        KwNode* node = &m->nodes[new_idx];

        KwNodeOpt new_opt = {
            .ch = kw[0],
            .is_term = false,
            .val = new_idx,
        };
        arena_array_push(&m->arena,
            (void**)&node->opts,
            &node->opts_cap,
            &node->opts_size,
            &new_opt,
            sizeof(new_opt));

        opt = &node->opts[node->opts_size - 1];
    }
    if (!kw[1])
        return opt;

    KwNode* next_node = &m->nodes[opt->val];
    return kwmatcher_create_path(m, next_node, &kw[1]);
}

static void kwmatcher_add(KwMatcher* m, const char* kw, TokTy tokty)
{
    KwNodeOpt* opt = kwmatcher_create_path(m, m->root, kw);
    opt->is_term = true;
    opt->val = tokty;
}

static KwNodeOpt* kwmatcher_find_node(
    KwMatcher* m, KwNode* node, const char* kw)
{
    assert(kw[0]);
    KwNodeOpt* opt = nullptr;
    for (size_t i = 0; i < node->opts_size; ++i) {
        if (node->opts[i].ch == kw[0]) {
            opt = &node->opts[i];
        }
    }
    if (!opt)
        return nullptr;
    if (!kw[1])
        return opt;
    return kwmatcher_find_node(m, &m->nodes[opt->val], &kw[1]);
}

static TokTy kwmatcher_match(KwMatcher* m, const char* kw)
{
    KwNodeOpt* opt = kwmatcher_find_node(m, m->root, kw);
    if (!opt->is_term)
        return Tok_Ident;
    return opt->val;
}

static void kwmatcher_init(KwMatcher* m, TokStream* tokstream)
{
    *m = (KwMatcher) {
        .tokstream = tokstream,
        .arena = { },
        ARRAY_FIELD_INIT(nodes),
        .root = nullptr,
    };

    arena_init(&m->arena, 512);

    arena_array_push(&m->arena,
        (void**)&(m)->nodes,
        &(m)->nodes_cap,
        &(m)->nodes_size,
        (&(KwNode) { .idx = 0 }),
        sizeof(*(&(KwNode) { .idx = 0 })));

    m->root = &m->nodes[0];

    struct {
        TokTy ty;
        const char* kw;
    } kws[] = {
        { Tok_KwInt, "int" },
        { Tok_KwChar, "char" },
        { Tok_KwConst, "const" },
    };

    for (size_t i = 0; i < sizeof(kws) / sizeof(kws[0]); ++i) {
        kwmatcher_add(m, kws[i].kw, kws[i].ty);
    }
}

static void kwmatcher_deinit(KwMatcher* m)
{
    arena_deinit(&m->arena);
}

static Tok kwmatcher_next(KwMatcher* m)
{
    Tok tok = tokstream_next(m->tokstream);
    if (tok.ty != Tok_Ident)
        return tok;
    TokTy ty = kwmatcher_match(m, tok.text);
    if (ty != Tok_Ident)
        printf("info: matched '%s' on line %d\n", tok.text, tok.loc.line);
    tok.ty = ty;
    return tok;
}

struct Parser {
    KwMatcher kwmatcher;
    Arena* arena;
    Arena sym_arena;

    Syms* ord_ns;
    Syms* label_ns;
    Syms* tag_ns;

    Tok curr;
};

Parser* parser_new(TokStream* tokstream, Arena* arena)
{
    Parser* p = malloc(sizeof(Parser));

    *p = (Parser) {
        .kwmatcher = { },
        .arena = arena,
        .sym_arena = { },

        .ord_ns = syms_new(),
        .label_ns = syms_new(),
        .tag_ns = syms_new(),

        .curr = { },
    };

    kwmatcher_init(&p->kwmatcher, tokstream);
    p->curr = kwmatcher_next(&p->kwmatcher);

    return p;
}

void parser_free(Parser* p)
{
    kwmatcher_deinit(&p->kwmatcher);
    free(p);
}

static void step(Parser* p)
{
    p->curr = kwmatcher_next(&p->kwmatcher);
}

#define DEFINE_NEW(TYPE, FN_NAME)                                              \
    [[maybe_unused]] static TYPE* new_##FN_NAME(Parser* p, TYPE init)          \
    {                                                                          \
        TYPE* val = arena_alloc(p->arena, sizeof(TYPE));                       \
        *val = init;                                                           \
        return val;                                                            \
    }

DEFINE_NEW(AstStmt, stmt)
DEFINE_NEW(AstDecl, decl)
DEFINE_NEW(AstDeclar, declar)

static bool is_decl_start_kw(TokTy ty)
{
    switch (ty) {
        case Tok_KwInt:
        case Tok_KwChar:
        case Tok_KwConst:
            return true;
        default:
            return false;
    }
}

static bool is_defty_ident(Parser* p, Tok tok)
{
    if (tok.ty != Tok_Ident)
        return false;
    Sym* sym = syms_get(p->ord_ns, tok.text);
    return sym && sym->ty == Sym_Typedef;
}

AstStmt* parser_parse_stmt(Parser* p)
{
    Loc loc = p->curr.loc;

    if (is_decl_start_kw(p->curr.ty) || is_defty_ident(p, p->curr)) {
        AstDecl* decl = parser_parse_decl(p);

        return new_stmt(p,
            (AstStmt) {
                .ty = AstStmt_Decl,
                .loc = loc,
                .decl = decl,
            });
    }

    AstExpr* expr = parser_parse_expr(p);

    return new_stmt(p,
        (AstStmt) {
            .ty = AstStmt_Expr,
            .loc = loc,
            .expr = expr,
        });
}

static AstDeclar* parse_declar(Parser* p);
static AstSpec parse_decl_spec(Parser* p);

static AstStmtBlock* parse_stmt_block(Parser* p);

AstDecl* parser_parse_decl(Parser* p)
{
    Loc loc = p->curr.loc;

    AstSpec spec = parse_decl_spec(p);

    ARRAY_VAR(AstDeclar*, declars);

    AstDeclar* first = parse_declar(p);

    array_push((void**)&declars,
        &declars_cap,
        &declars_size,
        (&first),
        sizeof(*(&first)));

    if (first->ty == AstDeclar_Function && p->curr.ty == '{') {
        AstStmtBlock* block = parse_stmt_block(p);
        return new_decl(p,
            (AstDecl) {
                .spec = spec,
                .declars = declars,
                .declars_cap = declars_cap,
                .declars_size = declars_size,
                .function_body = block,
                .loc = loc,
            });
    }

    while (p->curr.ty == ',') {
        step(p);

        AstDeclar* declar = parse_declar(p);

        array_push((void**)&declars,
            &declars_cap,
            &declars_size,
            (&declar),
            sizeof(*(&declar)));
    }

    return new_decl(p,
        (AstDecl) {
            .spec = spec,
            .declars = declars,
            .declars_cap = declars_cap,
            .declars_size = declars_size,
            .function_body = nullptr,
            .loc = loc,
        });
}

static AstDeclar* parse_declar_l1(Parser* p);
static AstDeclar* parse_declar_l2(Parser* p);
static AstDeclar* parse_declar_l3(Parser* p);

static AstDeclar* parse_declar(Parser* p)
{
    AstDeclar* declar = parse_declar_l1(p);

    if (p->curr.ty == '=') {
        step(p);
        AstExpr* expr = parser_parse_expr(p);
        declar->init_expr = expr;
    }

    return declar;
}

static AstDeclar* parse_declar_l1(Parser* p)
{
    Loc loc = p->curr.loc;

    if (p->curr.ty == '*') {
        step(p);

        AstSpec spec = parse_decl_spec(p);

        AstDeclar* declar = parse_declar_l1(p);
        return new_declar(p,
            (AstDeclar) {
                .ty = AstDeclar_Pointer,
                .loc = loc,
                .pointer = {
                    .declar = declar,
                    .spec = spec,
                },
            });
    }

    return parse_declar_l2(p);
}

static AstDeclar* parse_declar_l2(Parser* p)
{
    Loc loc = p->curr.loc;

    AstDeclar* declar = parse_declar_l3(p);

    while (p->curr.ty == '[') {
        step(p);
        AstExpr* expr = nullptr;
        if (p->curr.ty != ']') {
            expr = parser_parse_expr(p);
        }
        if (p->curr.ty != ']') {
            fprintf(
                stderr, "error: expected ']' on line %d\n", p->curr.loc.line);
            exit(1);
        }
        step(p);
        declar = new_declar(p,
            (AstDeclar) {
                .ty = AstDeclar_Array,
                .loc = loc,
                .array = {
                    .declar = declar,
                    .expr = expr,
                },
            });
    }

    return declar;
}

static AstDeclar* parse_declar_l3(Parser* p)
{
    Loc loc = p->curr.loc;

    if (p->curr.ty == Tok_Ident) {
        char* ident = arena_strdup(p->arena, p->curr.text);
        step(p);
        return new_declar(p,
            (AstDeclar) {
                .ty = AstDeclar_Array,
                .loc = loc,
                .ident = ident,
            });
    }

    if (p->curr.ty == '(') {
        step(p);

        AstDeclar* declar = parse_declar_l1(p);

        if (p->curr.ty != ')') {
            fprintf(
                stderr, "error: expected ')' on line %d\n", p->curr.loc.line);
            exit(1);
        }
        step(p);

        if (p->curr.ty != '(') {
            return declar;
        }
        step(p);

        ARRAY_VAR(AstDecl, params);

        while (true) {
            AstDecl* d = parser_parse_decl(p);
            arena_array_push(p->arena,
                (void**)&params,
                &params_cap,
                &params_size,
                (&d),
                sizeof(*(&d)));

            if (p->curr.ty != ',')
                break;
            step(p);
        }

        if (p->curr.ty != '(') {
            fprintf(
                stderr, "error: expected '(' on line %d\n", p->curr.loc.line);
            exit(1);
        }
        step(p);

        return new_declar(p,
            (AstDeclar) {
                .ty = AstDeclar_Function,
                .loc = loc,
                .function = {
                    .params = params,
                    .params_cap = params_cap,
                    .params_size = params_size,
                },
            });
    }

    fprintf(
        stderr, "error: expected declarator on line %d\n", p->curr.loc.line);
    exit(1);
}

static AstSpec parse_decl_spec(Parser* p)
{
    AstSpec spec = 0;

    while (true) {
        switch (p->curr.ty) {
            case Tok_Ident: {
                Sym* sym = syms_get(p->ord_ns, p->curr.text);
                if (sym->ty != Sym_Typedef)
                    break;
                spec |= AstSpec_DefTy;
                step(p);
                continue;
            }
            case Tok_KwInt:
                spec |= AstSpec_Int;
                step(p);
                continue;
            case Tok_KwChar:
                spec |= AstSpec_Char;
                step(p);
                continue;
            case Tok_KwConst:
                spec |= AstSpec_Const;
                step(p);
                continue;
            default:
        }

        break;
    }

    return spec;
}

static AstStmtBlock* parse_stmt_block(Parser* p)
{
    Loc loc = p->curr.loc;
}

AstExpr* parser_parse_expr(Parser* p)
{
}