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

struct MacroDef {
    Tok* toks;
    size_t toks_cap;
    size_t toks_size;

    char** params;
    size_t params_size;
};

struct MacroInv {
    MacroDef* def;
    size_t tok_idx;
    Hashmap arg_map;
};

typedef struct {
    Tok* toks;
    size_t toks_size;
} MacroArg;

void preproc_init(Preproc* p, Lexer* lexer)
{
    Tok curr = lexer_next(lexer);

    *p = (Preproc) {
        .lexer = lexer,
        .curr = curr,

        .defs = nullptr,
        .defs_cap = 0,
        .defs_size = 0,
        .def_map = { },

        .arg_arena = { },
        .macro_invs = nullptr,
        .macro_invs_cap = 0,
        .macro_invs_size = 0,
    };

    hashmap_init(&p->def_map);
    arena_init(&p->arg_arena, 512);
}

void preproc_deinit(Preproc* p)
{
    if (p->defs) {
        for (size_t i = 0; i < p->defs_size; ++i) {
            free(p->defs[i].toks);
            if (p->defs[i].params) {
                free(p->defs[i].params);
            }
        }
        free(p->defs);
    }
    hashmap_deinit(&p->def_map);
    arena_deinit(&p->arg_arena);
}

static void step(Preproc* p)
{
    p->curr = lexer_next(p->lexer);
}

static Tok p_define(Preproc* p)
{
    step(p);
    if (p->curr.ty != Tok_Ident) {
        fprintf(
            stderr, "error: expected identifier on line %d\n", p->curr.line);
        exit(1);
    }
    Tok ident_tok = p->curr;

    p->lexer->lex_whitespace = true;
    step(p);
    p->lexer->lex_whitespace = false;

    MacroDef def = { };

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

        char** params = nullptr;
        size_t params_cap = 0;
        size_t params_size = 0;

        if (p->curr.ty != Tok_Eof && p->curr.ty != ')') {
            if (p->curr.ty != Tok_Ident) {
                fprintf(stderr,
                    "error: expected identifier on line %d\n",
                    p->curr.line);
                exit(1);
            }
            array_push((void**)&params,
                &params_cap,
                &params_size,
                p->curr.text,
                sizeof(p->curr.text));

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

                if (p->curr.ty != Tok_Ident) {
                    fprintf(stderr,
                        "error: expected identifier on line %d\n",
                        p->curr.line);
                    exit(1);
                }
                array_push((void**)&params,
                    &params_cap,
                    &params_size,
                    p->curr.text,
                    sizeof(p->curr.text));
            }
        }

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

        def.params = params;
        def.params_size = params_size;
    }

    while (p->curr.ty != Tok_Eof && p->curr.ty != Tok_Eol) {
        array_push((void**)&def.toks,
            &def.toks_cap,
            &def.toks_size,
            &p->curr,
            sizeof(p->curr));
        step(p);
    }

    array_push(
        (void**)&p->defs, &p->defs_cap, &p->defs_size, &def, sizeof(def));

    MacroDef* deftoks = &p->defs[p->defs_size - 1];
    hashmap_set(&p->def_map, ident_tok.text, deftoks);

    return p->curr;
}

static Tok p_undef(Preproc* p)
{
    step(p);

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

    if (p->curr.ty != Tok_Eof && p->curr.ty != Tok_Eol) {
        fprintf(stderr,
            "error: unexpected extra tokens on line %d\n",
            p->curr.line);
        exit(1);
    }

    hashmap_set(&p->def_map, ident_tok.text, nullptr);

    return p->curr;
}

static Tok p_include(Preproc* p)
{
    p->lexer->lex_headername = true;
    step(p);
    p->lexer->lex_headername = false;

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

    assert(false && "not implemented");
}

static Tok p_paste_macro_body(Preproc* p)
{
    MacroInv* inv = &p->macro_invs[p->macro_invs_size - 1];

    if (inv->tok_idx >= inv->def->toks_size) {
        p->macro_invs_size -= 1;
        return preproc_next_with_newline(p);
    }

    Tok tok = inv->def->toks[inv->tok_idx];
    inv->tok_idx += 1;

    return tok;
}

static void collect_arg_toks(Tok** toks, size_t* toks_size)
{
}

struct arg_def {
    Tok* toks;
    size_t toks_size;
};

static Tok p_detect_macro(Preproc* p, Tok head)
{
    MacroDef* def = hashmap_get(&p->def_map, head.text);
    if (!def)
        return head;

    MacroInv inv = {
        .def = def,
        .tok_idx = 0,
        .arg_map = { },
    };
    hashmap_init(&inv.arg_map);

    array_push(
        (void**)&p->macro_invs, &p->defs_cap, &p->defs_size, &inv, sizeof(inv));

    if (!def->params)
        return preproc_next_with_newline(p);

    if (p->curr.ty != '(')
        return head;

    step(p);

    struct arg_def* arg_defs;
    size_t arg_defs_cap = 0;
    size_t arg_defs_size = 0;

    size_t arg_idx = 0;

    bool first = true;

    while ((first && (p->curr.ty != Tok_Eof && p->curr.ty != ')'))
        || p->curr.ty != ',') {
        first = false;

        if (arg_idx + 1 >= def->params_size) {
            fprintf(stderr,
                "error: too many arguments for macro line %d\n",
                p->curr.line);
            exit(1);
        }

        Tok* toks;
        size_t toks_size;
        collect_arg_toks(&toks, &toks_size);

        arena_array_push(&p->arg_arena,
            (void**)&arg_defs,
            &arg_defs_cap,
            &arg_defs_size,
            &(struct arg_def) { toks, toks_size },
            sizeof(struct arg_def));
        hashmap_set(
            &inv.arg_map, def->params[arg_idx], &arg_defs[arg_defs_size - 1]);

        arg_idx += 1;
    }

    return preproc_next_with_newline(p);
}

static Tok p_interpret_directive(Preproc* p)
{
    Tok kw = p->curr;

    if (strcmp(kw.text, "define") == 0) {
        return p_define(p);
    } else if (strcmp(kw.text, "undef") == 0) {
        return p_undef(p);
    } else if (strcmp(kw.text, "include") == 0) {
        return p_include(p);
    } else if (strcmp(kw.text, "if") == 0) {
        assert(false && "not implemented");
    } else if (strcmp(kw.text, "ifdef") == 0) {
        assert(false && "not implemented");
    } else if (strcmp(kw.text, "ifndef") == 0) {
        assert(false && "not implemented");
    } else if (strcmp(kw.text, "else") == 0) {
        assert(false && "not implemented");
    } else if (strcmp(kw.text, "elif") == 0) {
        assert(false && "not implemented");
    } else if (strcmp(kw.text, "elifdef") == 0) {
        assert(false && "not implemented");
    } else if (strcmp(kw.text, "elifndef") == 0) {
        assert(false && "not implemented");
    } else if (strcmp(kw.text, "endif") == 0) {
        assert(false && "not implemented");
    } else if (strcmp(kw.text, "line") == 0) {
        assert(false && "not implemented");
    } else if (strcmp(kw.text, "embed") == 0) {
        assert(false && "not implemented");
    } else if (strcmp(kw.text, "error") == 0) {
        assert(false && "not implemented");
    } else if (strcmp(kw.text, "warning") == 0) {
        assert(false && "not implemented");
    } else if (strcmp(kw.text, "pragma") == 0) {
        fprintf(stderr,
            "error: '%s' is not a preprocessor instruction on line %d\n",
            kw.text,
            kw.line);
        exit(1);
    }

    assert(false && "unreachable");
}

static Tok replace_candidate_tok(Preproc* p, Tok tok, size_t inv_idx)
{
    if (inv_idx == 0)
        return tok;
    MacroInv* inv = &p->macro_invs[inv_idx];
    if (tok.ty != Tok_Ident)
        return tok;
    MacroArg* arg = hashmap_get(&inv->arg_map, tok.text);
    if (!arg)
        return replace_candidate_tok(p, tok, inv_idx - 1);
}

static Tok expand_tok(Preproc* p, Tok tok)
{
    if (p->macro_invs_size == 0)
        return tok;

    MacroInv* inv = &p->macro_invs[p->macro_invs_size - 1];

    if (inv->tok_idx >= inv->def->toks_size) {
        p->macro_invs_size -= 1;
        return expand_tok(p, tok);
    }
}

Tok preproc_next_with_newline(Preproc* p)
{
    Tok tok = p->curr;
    step(p);
    if (tok.ty == '#')
        return p_interpret_directive(p);
    return expand_tok(p, tok);
}

Tok preproc_next(Preproc* p)
{
    Tok tok;
    while ((tok = preproc_next_with_newline(p)).ty == Tok_Eol
        || tok.ty == Tok_Whitespace) { }
    return tok;
}