#include <assert.h>
#include <string.h>

#include "ast.h"
#include "io.h"
#include "module.h"
#include "resolver.h"
#include "scanner.h"
#include "symtab.h"

/* Symbol storage across all scopes. */
static symbol_t SYMBOLS[MAX_SYMBOLS] = { 0 };
static usize    NSYMBOLS             = 0;

/* Scope storage across all modules/functions. */
static scope_t SCOPES[MAX_SCOPES] = { 0 };
static usize   NSCOPES            = 0;

static symkind_t symbol_entry(node_t *n) {
    switch (n->cls) {
    case NODE_IDENT:
    case NODE_VAR:
    case NODE_STATIC:
    case NODE_PARAM:
    case NODE_UNION_VARIANT: /* This should use SYM_VARIANT */
    case NODE_RECORD_FIELD:
        return SYM_VARIABLE;
    case NODE_CONST:
        return SYM_CONSTANT;
    case NODE_FN:
        return SYM_FUNCTION;
    case NODE_UNION:
    case NODE_RECORD:
    case NODE_PTR:
        return SYM_TYPE;
    case NODE_MOD:
    case NODE_USE:
        return SYM_MODULE;
    case NODE_REF:
        return symbol_entry(n->val.ref.target);
    case NODE_RETURN:
    case NODE_BLOCK:
    case NODE_LOOP:
    case NODE_WHILE:
    case NODE_WHILE_LET:
    case NODE_FOR:
    case NODE_IF:
    case NODE_IF_LET:
    case NODE_IF_CASE:
    case NODE_GUARD_CASE:
    case NODE_GUARD_LET:
    case NODE_BINOP:
    case NODE_UNOP:
    case NODE_TYPE:
    case NODE_NUMBER:
    case NODE_CHAR:
    case NODE_NIL:
    case NODE_UNDEF:
    case NODE_STRING:
    case NODE_BOOL:
    case NODE_ASSIGN:
    case NODE_CALL:
    case NODE_CALL_ARG:
    case NODE_BUILTIN:
    case NODE_ATTRIBUTE:
    case NODE_BREAK:
    case NODE_RECORD_TYPE:
    case NODE_RECORD_LIT:
    case NODE_ARRAY_LIT:
    case NODE_ARRAY_REPEAT_LIT:
    case NODE_ACCESS:
    case NODE_RECORD_LIT_FIELD:
    case NODE_EXPR_STMT:
    case NODE_ARRAY_INDEX:
    case NODE_RANGE:
    case NODE_MATCH_CASE:
    case NODE_MATCH:
    case NODE_SCOPE:
    case NODE_MOD_BODY:
    case NODE_AS:
    case NODE_PLACEHOLDER:
    case NODE_ALIGN:
    case NODE_THROW:
    case NODE_TRY:
    case NODE_CATCH:
    case NODE_PANIC:
    case NODE_SUPER:
        break;
    }
    bail("node of class %d cannot have a symbol table entry", n->cls);
}

/* Allocate a scope. */
scope_t *symtab_scope(scope_t *parent, module_t *mod) {
    if (NSCOPES >= MAX_SCOPES) {
        bail("scope overflow: too many scopes");
        return NULL;
    }
    scope_t *slot = &SCOPES[NSCOPES++];

    *slot = (scope_t){
        .mod      = mod,
        .parent   = parent,
        .symbols  = { 0 },
        .nsymbols = 0,
    };
    return slot;
}

/* Search for a symbol in the given scope only. */
symbol_t *symtab_scope_lookup(
    scope_t *s, const char *name, u16 length, symkind_t kind
) {
    for (usize i = 0; i < s->nsymbols; i++) {
        symbol_t *sym = s->symbols[i];

        if ((kind == SYM_ANY || sym->kind == kind) && sym->length == length &&
            memcmp(sym->name, name, length) == 0) {
            return sym;
        }
    }
    return NULL;
}

/* Search for a symbol from the current to the top-level scope. */
symbol_t *symtab_lookup(
    scope_t *s, const char *name, u16 length, symkind_t kind
) {
    for (scope_t *scope = s; scope != NULL; scope = scope->parent) {
        symbol_t *sym = NULL;

        if ((sym = symtab_scope_lookup(scope, name, length, kind)))
            return sym;
    }
    return NULL;
}

/* Add a symbol to the current scope. */
bool symtab_add_ident(scope_t *s, node_t *ident, node_t *n) {
    symkind_t   kind   = symbol_entry(n);
    const char *name   = ident->val.ident.name;
    u16         length = ident->val.ident.length;

    if (symtab_scope_lookup(s, name, length, kind))
        return false; /* Variable already defined in this scope. */

    symbol_t *sym = alloc_symbol((symbol_t){
        .name   = name,
        .length = length,
        .node   = n,
        .kind   = kind,
        .scope  = s,
    });

    /* Nb. static variables are never exposed. */
    if (kind == SYM_FUNCTION || kind == SYM_CONSTANT) {
        /* Copy qualified name for top-level symbols */

        module_path(sym->qualified, s->mod->qualified);
        module_qualify_str(sym->qualified, sym->name, sym->length);
    }
    assert(s->nsymbols < MAX_SCOPE_SYMBOLS);
    s->symbols[s->nsymbols++] = sym;
    n->sym                    = sym;

    return true;
}

/* Add a symbol to the current scope. */
bool symtab_insert(scope_t *s, const char *name, u16 length, node_t *n) {
    symkind_t kind = symbol_entry(n);

    if (symtab_lookup(s, name, length, kind)) {
        return false;
    }
    symbol_t *sym = alloc_symbol((symbol_t){
        .name   = name,
        .length = length,
        .node   = n,
        .kind   = kind,
    });

    assert(s->nsymbols < MAX_SCOPE_SYMBOLS);
    s->symbols[s->nsymbols++] = sym;
    n->sym                    = sym;

    return true;
}

/* Allocate a symbol. */
symbol_t *alloc_symbol(symbol_t sym) {
    if (NSYMBOLS >= MAX_SYMBOLS) {
        bail("symbol table overflow: too many symbols");
        return NULL;
    }
    symbol_t *slot = &SYMBOLS[NSYMBOLS++];
    *slot          = sym;

    strncpy(slot->qualified, slot->name, slot->length);

    return slot;
}

/* Add an imported symbol as an alias in the current scope. */
bool symtab_add_alias(scope_t *s, node_t *ident, symbol_t *original) {
    /* Check for conflicts */
    if (symtab_scope_lookup(
            s, ident->val.ident.name, ident->val.ident.length, original->kind
        )) {
        return false;
    }
    /* Create alias that points to the original symbol instead of copying it */
    assert(s->nsymbols < MAX_SCOPE_SYMBOLS);
    s->symbols[s->nsymbols++] = original;
    ident->sym                = original;

    return true;
}

/* Add a symbol directly to a scope. */
void symtab_add_symbol(scope_t *s, symbol_t *sym) {
    assert(s->nsymbols < MAX_SCOPE_SYMBOLS);
    s->symbols[s->nsymbols++] = sym;
}