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

typedef struct {
    ARRAY_FIELD(Hashmap, maps);
} Syms;

static void syms_push_layer(Syms* s);

static void syms_init(Syms* s)
{
    *s = (Syms) {
        ARRAY_FIELD_INIT(maps),
    };
    syms_push_layer(s);
}

static void syms_deinit(Syms* s)
{
    if (s->maps) {
        for (size_t i = 0; i < s->maps_size; ++i) {
            hashmap_deinit(&s->maps[i]);
        }
        free(s->maps);
    }
}

static void syms_push_layer(Syms* s)
{
    Hashmap m = { };
    hashmap_init(&m, sizeof(size_t));
    array_push((void**)&s->maps, &s->maps_cap, &s->maps_size, &m, sizeof(m));
}

static void syms_pop_layer(Syms* s)
{
    assert(s->maps_size > 0);

    hashmap_deinit(&s->maps[s->maps_size - 1]);
    s->maps_size -= 1;
}

static Hashmap* syms_current(Syms* s)
{
    assert(s->maps_size > 0);
    return &s->maps[s->maps_size - 1];
}

static void syms_define(Syms* s, const char* ident, size_t def_idx)
{
    Hashmap* m = syms_current(s);
    hashmap_set(m, ident, &def_idx);
}

static size_t* syms_get(Syms* s, const char* ident)
{
    for (size_t i = 0; i < s->maps_size; ++i) {
        Hashmap* m = &s->maps[s->maps_size - i - 1];
        size_t* def_idx = hashmap_get(m, ident);
        if (def_idx)
            return def_idx;
    }
    return nullptr;
}

typedef struct {
    size_t conds;
    size_t depth;
} CondStack;

void condstack_push(CondStack* s, bool cond)
{
    if (cond) {
        s->conds |= 1ul << s->depth;
    } else {
        s->conds &= ~(1ul << s->depth);
    }
    s->depth += 1;
}

bool condstack_pop(CondStack* s)
{
    assert(s->depth > 0);
    s->depth -= 1;
    return (s->conds >> s->depth) & 1;
}

typedef struct {
    char* filename;
    char* text;
    SourceIter source_iter;
    LslIter lsl_iter;
    Lexer lexer;
} File;

typedef struct {
    Tok* toks;
    size_t toks_size;

    char** params;
    size_t params_size;
} Def;

struct Preproc {
    Arena* arena;
    ARRAY_FIELD(File, files);

    ARRAY_FIELD(Tok, tok_stack);
    Tok curr;

    ARRAY_FIELD(Def, defs);
    Syms syms;

    CondStack condstack;

    bool next_can_be_directive;
};

void preproc_init(Preproc** p, char* filename, char* text, Arena* arena)
{
    *p = malloc(sizeof(Preproc));

    **p = (Preproc) {
        .arena = arena,
        ARRAY_FIELD_INIT(files),

        ARRAY_FIELD_INIT(tok_stack),
        .curr = { },

        ARRAY_FIELD_INIT(defs),
        .syms = { },

        .condstack = { },

        .next_can_be_directive = true,
    };

    ARRAY_FIELD_PUSH(*p,
        files,
        (&(File) {
            .filename = filename,
            .text = text,
            .source_iter = { },
            .lsl_iter = { },
            .lexer = { },
        }));

    File* file = &(*p)->files[0];
    source_iter_init(&file->source_iter, text);
    lsl_iter_init(&file->lsl_iter, &file->source_iter);
    lexer_init(&file->lexer, &file->lsl_iter, (*p)->arena);
    (*p)->curr = lexer_next(&file->lexer);

    // null-definition at idx 0
    ARRAY_FIELD_PUSH(*p, defs, &(Def) { });

    syms_init(&(*p)->syms);
}

void preproc_deinit(Preproc* p)
{
    for (size_t i = 0; i < p->files_size; ++i) {
        free(p->files[i].filename);
        free(p->files[i].text);
    }
    free(p->files);
    if (p->tok_stack)
        free(p->tok_stack);
    if (p->defs) {
        for (size_t i = 0; i < p->defs_size; ++i) {
            if (p->defs[i].params)
                free(p->defs[i].params);
        }
        free(p->defs);
    }
    syms_deinit(&p->syms);
    free(p);
}

static void step(Preproc* p)
{
    p->next_can_be_directive = p->curr.ty == Tok_Eol;

    if (p->tok_stack_size == 0) {
        p->curr = lexer_next(&p->files[p->files_size - 1].lexer);

        if (p->curr.ty == Tok_Eof && p->files_size > 1) {
            File* curr_file = &p->files[p->files_size - 1];
            free(curr_file->filename);
            free(curr_file->text);

            p->files_size -= 1;
            step(p);
        }

        return;
    }
    p->curr = p->tok_stack[p->tok_stack_size - 1];
    p->tok_stack_size -= 1;

    if (p->curr.ty == Tok_EndOfMacroInvocation) {
        syms_pop_layer(&p->syms);
        step(p);
        return;
    }
}

static void collect_arg(Preproc* p, size_t def_idx, size_t arg_idx)
{
    if (arg_idx >= p->defs[def_idx].params_size) {
        fprintf(stderr,
            "error: too many arguments for macro on line %d\n",
            p->curr.line);
        exit(1);
    }

    ARRAY_VAR(Tok, toks);

    int depth = 0;

    while (p->curr.ty != Tok_Eof
        && (depth != 0 || (p->curr.ty != ',' && p->curr.ty != ')'))) {

        if (p->curr.ty == '(')
            depth += 1;
        if (p->curr.ty == ')') {
            if (depth == 0)
                continue;
            depth -= 1;
        }

        arena_array_push(p->arena,
            (void**)&toks,
            &toks_cap,
            &toks_size,
            (&p->curr),
            sizeof(*(&p->curr)));

        step(p);
    }

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

    ARRAY_FIELD_PUSH(p,
        defs,
        (&(Def) {
            .toks = toks,
            .toks_size = toks_size,
        }));

    syms_define(&p->syms, p->defs[def_idx].params[arg_idx], p->defs_size - 1);
}

#define A(B, C) B C

static void collect_args(Preproc* p, size_t def_idx)
{
    assert(p->curr.ty == '(');
    step(p);

    size_t idx = 0;

    if (p->curr.ty != Tok_Eof && p->curr.ty != ')') {
        collect_arg(p, def_idx, idx);
        idx += 1;
        while (p->curr.ty == ',') {
            step(p);
            collect_arg(p, def_idx, idx);
            idx += 1;
        }
    }
    if (p->curr.ty != ')') {
        fprintf(stderr, "error: expected ')' on line %d\n", p->curr.line);
        exit(1);
    }

    assert(idx >= p->defs[def_idx].params_size);
    if (idx > p->defs[def_idx].params_size) {
        fprintf(stderr,
            "error: too few arguments for macro on line %d\n",
            p->curr.line);
        exit(1);
    }
}

static Tok expand_tok(Preproc* p)
{
    Tok tok = p->curr;
    size_t* def_idx;
    if (tok.ty != Tok_Ident || !(def_idx = syms_get(&p->syms, tok.text))) {
        step(p);
        return tok;
    }

    // macro expension doesn't support recursion.
    // when seeing the same macro a second time
    // in the invocation tree, treat it as a normal token.
    // undefine it at this layer.
    syms_push_layer(&p->syms);
    syms_define(&p->syms, tok.text, 0);

    if (p->defs[*def_idx].params) {
        step(p);
        // function-like macros without arg-lists are ignored.
        if (p->curr.ty != '(')
            return tok;
        collect_args(p, *def_idx);
    }

    ARRAY_FIELD_PUSH(p, tok_stack, (&(Tok) { .ty = Tok_EndOfMacroInvocation }));
    for (size_t i = 0; i < p->defs[*def_idx].toks_size; ++i) {
        Def* def = &p->defs[*def_idx];
        Tok expanded_tok = def->toks[def->toks_size - i - 1];
        expanded_tok.line = tok.line;
        expanded_tok.col = tok.col;
        ARRAY_FIELD_PUSH(p, tok_stack, &expanded_tok);
    }

    // step into the new invocation
    step(p);

    return expand_tok(p);
}

static Tok do_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->files[p->files_size - 1].lexer.lex_whitespace = true;
    step(p);
    p->files[p->files_size - 1].lexer.lex_whitespace = false;

    Def def = { };

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

        ARRAY_VAR(char*, params);

        if (p->curr.ty != Tok_Eof && p->curr.ty != ')') {
            if (p->curr.ty != Tok_Ident)
                goto error_expected_identifier;

            ARRAY_VAR_PUSH(params, &p->curr.text);
            step(p);

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

                if (p->curr.ty != Tok_Ident)
                    goto error_expected_identifier;

                ARRAY_VAR_PUSH(params, &p->curr.text);
                step(p);
            }
        }

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

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

    ARRAY_VAR(Tok, toks);

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

    def.toks = toks;
    def.toks_size = toks_size;

    ARRAY_FIELD_PUSH(p, defs, &def);
    syms_define(&p->syms, ident_tok.text, p->defs_size - 1);

    return p->curr;

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

static Tok do_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);
    }

    syms_define(&p->syms, ident_tok.text, 0);

    return p->curr;
}

Tok preproc_next(Preproc* p);

static Tok do_include(Preproc* p)
{
    File* curr_file = &p->files[p->files_size - 1];

    curr_file->lexer.lex_headername = true;
    step(p);
    curr_file->lexer.lex_headername = false;

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

    Tok headername_tok = p->curr;
    step(p);

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

    char* rel_path
        = strndup(headername_tok.text + 1, strlen(headername_tok.text) - 2);

    char* last_delim = strrchr(curr_file->filename, '/');

    char* filename;
    if (last_delim) {
        size_t path_len = last_delim - curr_file->filename;

        filename = calloc(path_len + 1 + strlen(rel_path) + 1, 1);

        strncat(filename, curr_file->filename, path_len);
        strcat(filename, "/");
        strcat(filename, rel_path);
    } else {
        filename = strdup(rel_path);
    }
    free(rel_path);

    char* text = read_textfile(filename);

    ARRAY_FIELD_PUSH(p,
        files,
        (&(File) {
            .filename = filename,
            .text = text,
            .source_iter = { },
            .lsl_iter = { },
            .lexer = { },
        }));

    File* file = &p->files[p->files_size - 1];
    source_iter_init(&file->source_iter, text);
    lsl_iter_init(&file->lsl_iter, &file->source_iter);
    lexer_init(&file->lexer, &file->lsl_iter, p->arena);
    p->curr = lexer_next(&file->lexer);

    p->next_can_be_directive = true;

    return preproc_next(p);
}

static Tok interpret_directive(Preproc* p);

static Tok do_ifndef(Preproc* p)
{
    int original_line = p->curr.line;

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

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

    bool cond = syms_get(&p->syms, ident.text) == nullptr;

    condstack_push(&p->condstack, cond);

    if (!cond) {

        while (p->curr.ty != Tok_Eof) {
            while (p->curr.ty != Tok_Eof
                && (!p->next_can_be_directive || p->curr.ty != '#'))
                step(p);
            if (p->curr.ty == Tok_Eof)
                break;
            step(p);
            if (p->curr.ty != Tok_Ident)
                continue;
            if (strcmp(p->curr.text, "endif") == 0)
                return interpret_directive(p);
        }

        if (p->curr.ty != Tok_Eof) {
            fprintf(stderr,
                "error: unclosed '#ifndef' on line %d\n",
                original_line);
            exit(1);
        }
    }

    return p->curr;
}

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

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

    condstack_pop(&p->condstack);

    return p->curr;
}

static Tok interpret_directive(Preproc* p)
{
    if (p->curr.ty == Tok_Eol)
        return p->curr;

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

    Tok kw = p->curr;

    struct {
        const char* keyword;
        Tok (*fn)(Preproc* p);
    } directives[] = {
        { "define", do_define },
        { "undef", do_undef },
        { "include", do_include },
        { "ifndef", do_ifndef },
        { "endif", do_endif },

        { "if", nullptr },
        { "ifdef", nullptr },
        { "else", nullptr },
        { "elif", nullptr },
        { "elifdef", nullptr },
        { "elifndef", nullptr },
        { "line", nullptr },
        { "embed", nullptr },
        { "error", nullptr },
        { "warning", nullptr },
        { "pragma", nullptr },
    };

    for (size_t i = 0; i < sizeof(directives) / sizeof(directives[0]); ++i) {
        if (strcmp(kw.text, directives[i].keyword) == 0) {
            assert(directives[i].fn && "not implemented");

            return directives[i].fn(p);
        }
    }

    fprintf(stderr,
        "error: '%s' is not a preprocessor instruction on line %d\n",
        kw.text,
        kw.line);
    exit(1);
}

Tok preproc_next(Preproc* p)
{
    if (p->next_can_be_directive && p->curr.ty == '#') {
        step(p);
        return interpret_directive(p);
    }
    return expand_tok(p);
}

void tokstream_init(TokStream* s, Preproc* preproc, Arena* arena)
{
    *s = (TokStream) {
        .preproc = preproc,
        .arena = arena,
        .curr = { },
    };
    s->curr = preproc_next(s->preproc);
}

static Tok tokstream_next_filtered(TokStream* s)
{
    Tok tok;
    while ((tok = preproc_next(s->preproc)).ty == Tok_Eol
        || tok.ty == Tok_Whitespace) { }
    return tok;
}

Tok tokstream_next(TokStream* s)
{
    Tok curr = s->curr;
    Tok next = tokstream_next_filtered(s);

    if (curr.ty == Tok_Str && next.ty == Tok_Str) {
        char* text
            = arena_alloc(s->arena, strlen(curr.text) + strlen(next.text) + 1);
        text[0] = '\0';
        strcat(text, curr.text);
        strcat(text, next.text);
        s->curr = (Tok) { .ty = Tok_Str, text, curr.line, curr.col };
        return tokstream_next(s);
    }

    s->curr = next;
    return curr;
}