#include "parse.h"
#include "ast.h"
#include "collection.h"
#include "syms.h"
#include <assert.h>
#include <stddef.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.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);
} KwMatcher;

static void kw_tree_print(KwMatcher* m, KwNode* node, int depth)
{
    for (size_t i = 0; i < node->opts_size; ++i) {
        for (int i = 0; i < depth; ++i)
            printf(" ");
        printf("%c", node->opts[i].ch);
        if (node->opts[i].is_term) {
            printf(" = %s\n", tokty_str(node->opts[i].val));
        } else {
            printf("\n");
            kw_tree_print(m, &m->nodes[node->opts[i].val], depth + 1);
        }
    }
}

static KwNodeOpt* kwmatcher_create_path(
    KwMatcher* m, size_t node_idx, const char* kw)
{
    assert(kw[0]);
    KwNode* node = &m->nodes[node_idx];

    KwNodeOpt* opt = nullptr;
    for (size_t i = 0; i < node->opts_size; ++i) {
        if (node->opts[i].ch == kw[0]) {
            opt = &node->opts[i];
            break;
        }
    }
    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));
        // might have been relocated
        node = &m->nodes[node_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;

    return kwmatcher_create_path(m, opt->val, &kw[1]);
}

static void kwmatcher_add(KwMatcher* m, const char* kw, TokTy tokty)
{
    KwNodeOpt* opt = kwmatcher_create_path(m, 0, 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->nodes[0], kw);
    if (!opt || !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),
    };

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

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

    kw_tree_print(m, &m->nodes[0], 0);
}

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)
        goto exit;
    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;

exit:
    printf("  - %s(%s) %d:%d\n",
        tokty_str(tok.ty),
        tok.text,
        tok.loc.line,
        tok.loc.col);
    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)
DEFINE_NEW(AstExpr, expr)

DEFINE_NEW(AstStmtBlock, stmt_block)

static AstStmtBlock* parse_stmt_block(Parser* p);
static AstDecl* parse_decl(Parser* p);
static AstDeclar* parse_declar(Parser* p);
static AstSpec parse_decl_spec(Parser* p);
static AstExpr* parse_expr(Parser* p);

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

    ARRAY_VAR(AstStmt*, stmts);

    while (p->curr.ty != Tok_Eof && p->curr.ty != '}') {
        AstStmt* stmt = parser_parse_stmt(p);
        arena_array_push(p->arena,
            (void**)&stmts,
            &stmts_cap,
            &stmts_size,
            (&stmt),
            sizeof(*(&stmt)));
    }

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

    return new_stmt_block(p,
        (AstStmtBlock) {
            .stmts = stmts,
            .stmts_cap = stmts_cap,
            .stmts_size = stmts_size,
            .loc = loc,
        });
}

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

static AstStmt* 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 = parse_decl(p);

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

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

    AstExpr* expr = parse_expr(p);

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

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

    AstSpec spec = parse_decl_spec(p);

    ARRAY_VAR(AstDeclar*, declars);

    AstDeclar* first = parse_declar(p);

    ARENA_ARRAY_VAR_PUSH(p->arena, declars, &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_l4(Parser* p);

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

    if (p->curr.ty == '=') {
        step(p);
        AstExpr* expr = 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 = 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;

    AstDeclar* declar = parse_declar_l4(p);

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

        ARRAY_VAR(AstDecl*, params);

        while (true) {
            AstDecl* d = parse_decl(p);
            ARENA_ARRAY_VAR_PUSH(p->arena, params, &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);

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

    return declar;
}

static AstDeclar* parse_declar_l4(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);

        return declar;
    }

    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 || 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 AstExpr* parse_expr_1(Parser* p);
static AstExpr* parse_expr_2(Parser* p);

static AstExpr* parse_expr(Parser* p)
{
    return parse_expr_1(p);
}

static AstExpr* parse_expr_1(Parser* p)
{
    Loc loc = p->curr.loc;

    AstExpr* expr = parse_expr_2(p);

    while (true) {

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

            ARRAY_VAR(AstExpr*, args);

            if (p->curr.ty != ')') {
                do {
                    step(p);
                    AstExpr* arg = parse_expr(p);
                    ARRAY_VAR_PUSH(args, &arg);
                } while (p->curr.ty == ',');
            }

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

            expr = new_expr(p,
                (AstExpr) {
                    .ty = AstExpr_Call,
                    .loc = loc,
                    .call = {
                        .expr = expr,
                        .args = args,
                        .args_cap = args_cap,
                        .args_size = args_size,
                    },
                });
        }

        break;
    }

    return expr;
}

static size_t unescape_string(char* dest, const char* src)
{
    size_t dest_i = 0;
    size_t src_len = strlen(src);

    assert(src[0] == '"' && src[src_len - 1] == '"');

    src -= 1;
    src_len -= 2;

    for (size_t i = 0; i < src_len; ++i) {
        char ch = src[i];

        char dest_ch;
        if (ch == '\\') {
            i += 1;
            assert(i < src_len);
            ch = src[i];

            switch (ch) {
                case 't':
                    dest_ch = '\t';
                    break;
                case 'r':
                    dest_ch = '\r';
                    break;
                case 'n':
                    dest_ch = '\n';
                    break;
                case '0':
                    dest_ch = '\0';
                    break;
                default:
                    dest_ch = ch;
            }

        } else {
            dest_ch = ch;
        }
        if (dest)
            dest[dest_i] = dest_ch;
        dest_i += 1;
    }

    return dest_i;
}

static AstExpr* parse_expr_2(Parser* p)
{
    Loc loc = p->curr.loc;

    if (p->curr.ty == Tok_Str) {
        size_t unescaped_len = unescape_string(nullptr, p->curr.text);
        char* unescaped_str = arena_alloc(p->arena, unescaped_len + 1);
        unescape_string(unescaped_str, p->curr.text);

        step(p);

        return new_expr(p,
            (AstExpr) {
                .ty = AstExpr_Str,
                .loc = loc,
                .str_val = unescaped_str,
            });
    }

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

AstStmt* parser_parse_stmt(Parser* p)
{
    return parse_stmt(p);
}

bool parser_done(Parser* p)
{
    return p->curr.ty == Tok_Eof;
}