#include "parse.h"
#include "ast.h"
#include "collection.h"
#include "resolve.h"
#include "sema.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 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_FIELD_PUSH(&m->arena, m, nodes, &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_FIELD_PUSH(&m->arena, node, opts, &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_FIELD_PUSH(&m->arena, m, nodes, (&(KwNode) { .idx = 0 }));

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

    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)
        goto exit;
    TokTy ty = kwmatcher_match(m, tok.text);
    tok.ty = ty;

exit:
    printf("        - %-10s %-10s %d:%d\n",
        tokty_str(tok.ty),
        tok.text,
        tok.loc.line,
        tok.loc.col);
    return tok;
}

typedef enum : unsigned int {
    ParserFlags_ParsingTypedef = 1 << 0,
} ParserFlags;

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

    Syms* syms;

    Tok curr;

    ParserFlags flags;
};

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

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

        .syms = syms_new(),

        .curr = { },

        .flags = 0,
    };

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

    return p;
}

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

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

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

DEFINE_NEW(AstStmt, stmt, { sema_visit_stmt(p->sema, val); })
DEFINE_NEW(AstDecl, decl, { })
DEFINE_NEW(AstDeclar, declar, { })
DEFINE_NEW(AstExpr, expr, { sema_visit_expr(p->sema, val); })

DEFINE_NEW(AstStmtBlock, stmt_block, { })

static AstStmtBlock* parse_stmt_block(Parser* p);
static AstStmt* parse_stmt(Parser* p);
static AstDecl* parse_decl(Parser* p, bool is_param);
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 == '{');
    step(p);

    ARRAY_VAR(AstStmt*, stmts);

    while (p->curr.ty != Tok_Eof && p->curr.ty != '}') {
        AstStmt* stmt = parse_stmt(p);
        ARENA_ARRAY_VAR_PUSH(p->arena, stmts, &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;
    OldSym* sym = syms_get(p->syms, tok.text);
    return sym && sym->ty == OldSym_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, false);

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

    if (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_Expr,
            .loc = loc,
            .expr = expr,
        });
}

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

    AstSpec spec = parse_decl_spec(p);

    if (spec & AstSpec_Typedef) {
        p->flags |= ParserFlags_ParsingTypedef;
    }

    ARRAY_VAR(AstDeclar*, declars);

    AstDeclar* first = parse_declar(p);

    ARENA_ARRAY_VAR_PUSH(p->arena, declars, &first);

    AstStmtBlock* function_body = nullptr;
    bool multi_allowed = !is_param;

    if (!is_param && first->ty == AstDeclar_Function && p->curr.ty == '{') {

        sema_function_enter(p->sema, spec, declars[0]);
        function_body = parse_stmt_block(p);
        sema_function_leave(p->sema);

        multi_allowed = false;
    }

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

        AstDeclar* declar = parse_declar(p);

        ARENA_ARRAY_VAR_PUSH(p->arena, declars, &declar);
    }

    p->flags &= ~ParserFlags_ParsingTypedef;

    return new_decl(p,
        (AstDecl) {
            .spec = spec,
            .declars = declars,
            .declars_cap = declars_cap,
            .declars_size = declars_size,
            .function_body = function_body,
            .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);

        if (p->curr.ty == Tok_KwVoid) {
            step(p);
        } else if (p->curr.ty != Tok_Eof && p->curr.ty != ')') {
            AstDecl* d = parse_decl(p, true);
            ARENA_ARRAY_VAR_PUSH(p->arena, params, &d);

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

                d = parse_decl(p, true);
                ARENA_ARRAY_VAR_PUSH(p->arena, params, &d);
            }
        }

        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_Ident,
                .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: {
                OldSym* sym = syms_get(p->syms, p->curr.text);
                if (!sym || sym->ty != OldSym_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 == '(') {

            ARRAY_VAR(AstExpr*, args);

            if (p->curr.ty != ')') {
                do {
                    step(p);
                    AstExpr* arg = parse_expr(p);
                    ARENA_ARRAY_VAR_PUSH(p->arena, 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;
    }

    if (dest)
        dest[dest_i] = '\0';
    return dest_i;
}

static AstExpr* parse_expr_2(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_expr(p,
            (AstExpr) {
                .ty = AstExpr_Ident,
                .loc = loc,
                .ident = ident,
            });
    }

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