#include "resolve.h"
#include "ast.h"
#include "collection.h"
#include <assert.h>
#include <stdio.h>
#include <stdlib.h>

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

Syms* syms_new()
{
    Syms* s = malloc(sizeof(Syms));
    *s = (Syms) {
        ARRAY_FIELD_INIT(maps),
    };

    ARRAY_FIELD_PUSH(s, maps, &(Hashmap) { });
    hashmap_init(&s->maps[0], sizeof(Sym*));

    return s;
}

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

void syms_push_scope(Syms* s)
{
    ARRAY_FIELD_PUSH(s, maps, &(Hashmap) { });
    hashmap_init(&s->maps[s->maps_size - 1], sizeof(Sym*));
}

void syms_pop_scope(Syms* s)
{
    assert(s->maps_size > 0);
    hashmap_deinit(&s->maps[s->maps_size - 1]);
    s->maps_size -= 1;
}

void syms_define(Syms* s, const char* ident, Sym* sym)
{
    Hashmap* m = &s->maps[s->maps_size - 1];
    hashmap_set(m, ident, &sym);
}

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

static void print_sym(Sym* s)
{
    switch (s->ty) {
        case Sym_Typedef:
            printf("Typedef");
            break;
        case Sym_Declarator:
            printf("Declarator");
            break;
    }
}

struct SymPrinter {
    int depth;
};

static void print_sym_hashmap(void* self, const char* key, void* val)
{
    struct SymPrinter* p = self;
    Sym* sym = *(Sym**)val;

    printf("  ");
    for (int i = 0; i < p->depth; ++i)
        printf("  ");
    printf("'%s' -> ", key);
    print_sym(sym);
    printf("\n");
}

void syms_print(Syms* s)
{
    for (size_t i = 0; i < s->maps_size; ++i) {
        struct SymPrinter p = {
            .depth = i,
        };
        hashmap_for_each(&s->maps[i], print_sym_hashmap, &p);
    }
}

typedef enum : unsigned int {
    ResolverFlags_InTypedef,
} ResolverFlags;

struct Resolver {
    Syms* ord_syms;
    Syms* tag_syms;
    Syms* label_syms;

    Arena* arena;

    Keymap resol_map;

    AstStmt* current_decl_stmt;
    ResolverFlags flags;
};

Resolver* resolver_new(Arena* arena)
{
    Resolver* r = malloc(sizeof(Resolver));
    *r = (Resolver) {
        .ord_syms = syms_new(),
        .tag_syms = syms_new(),
        .label_syms = syms_new(),

        .arena = arena,

        .resol_map = { },

        .current_decl_stmt = nullptr,
        .flags = 0,
    };

    keymap_init(&r->resol_map, sizeof(Sym*));

    return r;
}

void resolver_free(Resolver* r)
{
    keymap_deinit(&r->resol_map);

    syms_free(r->ord_syms);
    syms_free(r->tag_syms);
    syms_free(r->label_syms);
    free(r);
}

static AstVisitorResult visit_stmt_block(AstVisitor* v, AstStmtBlock* node)
{
    Resolver* r = v->op;
    syms_push_scope(r->ord_syms);

    ast_stmt_block_visit_body(v, node);

    syms_pop_scope(r->ord_syms);
    return AstVisitor_Break;
}

static AstVisitorResult visit_stmt(AstVisitor* v, AstStmt* node)
{
    Resolver* r = v->op;

    if (node->ty == AstStmt_Decl) {
        r->current_decl_stmt = node;

        ast_stmt_visit_body(v, node);

        return AstVisitor_Break;
    }

    return AstVisitor_Continue;
}

static AstVisitorResult visit_decl(AstVisitor* v, AstDecl* node)
{
    Resolver* r = v->op;

    if (node->spec & AstSpec_Typedef) {
        r->flags |= ResolverFlags_InTypedef;

        ast_decl_visit_body(v, node);

        r->flags &= ~ResolverFlags_InTypedef;
        return AstVisitor_Break;
    }

    if (node->function_body) {

        assert(node->declars_size == 1
            && node->declars[0]->ty == AstDeclar_Function);

        AstDeclar* d = node->declars[0];

        ast_declar_visit(v, d->function.declar);

        syms_push_scope(r->ord_syms);

        for (size_t i = 0; i < d->function.params_size; ++i)
            ast_decl_visit(v, d->function.params[i]);

        ast_stmt_block_visit(v, node->function_body);

        syms_pop_scope(r->ord_syms);
        return AstVisitor_Break;
    }

    return AstVisitor_Continue;
}

static AstVisitorResult visit_declar(AstVisitor* v, AstDeclar* node)
{
    Resolver* r = v->op;

    if (node->ty == AstDeclar_Function) {
        ast_declar_visit(v, node->function.declar);

        syms_push_scope(r->ord_syms);

        for (size_t i = 0; i < node->function.params_size; ++i)
            ast_decl_visit(v, node->function.params[i]);

        syms_pop_scope(r->ord_syms);
        return AstVisitor_Break;
    }

    if (node->ty == AstDeclar_Ident) {
        assert(r->current_decl_stmt);

        Sym* sym = arena_alloc(r->arena, sizeof(Sym));
        *sym = (Sym) {
            .ty = Sym_Declarator,
            .declarator = {
                .stmt = r->current_decl_stmt,
                .declar = node,
            },
        };
        syms_define(r->ord_syms, node->ident, sym);
        return AstVisitor_Break;
    }

    return AstVisitor_Continue;
}

static AstVisitorResult visit_expr(AstVisitor* v, AstExpr* node)
{
    Resolver* r = v->op;

    if (node->ty == AstExpr_Ident) {

        Sym* sym = syms_get(r->ord_syms, node->ident);
        if (!sym) {
            fprintf(stderr,
                "error: use of undeclared '%s' on line %d\n",
                node->ident,
                node->loc.line);

            syms_print(r->ord_syms);

            exit(1);
        }
    }

    return AstVisitor_Continue;
}

void resolver_resolve_top_stmt(Resolver* r, AstStmt* stmt)
{
    if (stmt->ty != AstStmt_Decl) {
        fprintf(
            stderr, "error: expected declaration on line %d\n", stmt->loc.line);
        exit(1);
    }

    AstVisitor v = {
        .visit_stmt_block = visit_stmt_block,
        .visit_stmt = visit_stmt,
        .visit_decl = visit_decl,
        .visit_declar = visit_declar,
        .visit_expr = visit_expr,
        .op = r,
    };

    ast_stmt_visit(&v, stmt);
}