#include "parse.h"
#include "collection.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.line);
    tok.ty = ty;
    return tok;
}

struct Parser {
    KwMatcher kwmatcher;
    Arena* arena;

    Tok curr;
};

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

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

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