#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)
{
    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)
{
    Hashmap* m = &s->maps[s->maps_size - 1];
    Sym** res = hashmap_get(m, ident);
    return res ? *res : nullptr;
}

typedef enum : unsigned int {
    ResolverFlags_InTypedef,
} ResolverFlags;

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

    Arena* arena;

    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,

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

void resolver_free(Resolver* r)
{
    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;
    }

    return AstVisitor_Continue;
}

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

    if (node->ty == AstDeclar_Function) {
        syms_push_scope(r->ord_syms);

        ast_declar_visit_body(v, node);

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

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

    visit_stmt(&v, stmt);
}