#include "scope.h"
#include "ast.h"
#include <stdlib.h>
#include <stdio.h>
#include <string.h>
#define DEFAULT_SIZE 16

static void scope_init(struct scope* scope) {
    scope->types = calloc(DEFAULT_SIZE, sizeof(struct type_def*));
    scope->type_cap = DEFAULT_SIZE;
    scope->vars = calloc(DEFAULT_SIZE, sizeof(struct var_def*));
    scope->var_cap = DEFAULT_SIZE;
}

static void type_destroy(struct type_def* type) {
    free(type->key.name);
}

static void var_destroy(struct var_def* var) {
    free(var->name);
}

void scope_destroy(struct scope* scope) {
    for (unsigned long long i = 0; i < scope->type_cap; i++) {
        if (scope->types[i] != NULL) {
            type_destroy(scope->types[i]);
            free(scope->types[i]);
        }
    }
    free(scope->types);

    for (unsigned long long i = 0; i < scope->var_cap; i++) {
        if (scope->vars[i] != NULL) {
            var_destroy(scope->vars[i]);
            free(scope->vars[i]);
        }
    }
    free(scope->vars);
}

static inline unsigned long long advance_hash(
    unsigned long long hash,
    unsigned long long val,
    unsigned long long cap
) {
    return ((hash << 5) - hash + val) % cap;;
}

static unsigned long long hash_name(const char* name, unsigned long long cap) {
    unsigned long long hash = 0, i = 0;
    while (name[i] != 0) hash = advance_hash(hash, name[i++], cap);
    return hash;
}

unsigned long long hash_type_key(
    const struct type_key* key,
    unsigned long long cap
) {
    unsigned long long hash = hash_name(key->name, cap);
    hash = advance_hash(hash, key->how_long, cap);
    hash = advance_hash(hash, key->marked_signed, cap);
    hash = advance_hash(hash, key->marked_unsigned, cap);
    return hash;
}

static bool type_keys_equal(
    const struct type_key* a,
    const struct type_key* b
) {
    return a->how_long == b->how_long
        && a->marked_signed == b->marked_signed
        && a->marked_unsigned == b->marked_unsigned
        && strcmp(a->name, b->name) == 0;
}

static struct type_def** type_cell(
    const struct scope* scope,
    const struct type_key* key
) {
    unsigned long long orig_idx = hash_type_key(key, scope->type_cap);
    unsigned long long idx = orig_idx;

    do {
        if (scope->types[idx] == NULL
                || type_keys_equal(key, &scope->types[idx]->key))
            return &scope->types[idx];
    } while ((idx = (idx + 1) % scope->type_cap) != orig_idx);
    return NULL;
}

static void rehash_types(struct scope* scope) {
    struct type_def** old_types = scope->types;
    unsigned long long old_cap = scope->type_cap;
    scope->type_cap *= 2;
    scope->types = calloc(scope->type_cap, sizeof(struct type_def*));
    if (scope->types == NULL) {
        fprintf(stderr, "ccc: out of memory\n");
        exit(1);
    }

    for (unsigned long long i = 0; i < old_cap; i++) {
        if (old_types[i] == NULL) continue;

        struct type_def** cell = type_cell(scope, &old_types[i]->key);
        if (cell == NULL) {
            fprintf(stderr, "ccc: types rehash failed, likely a bug\n");
            exit(1);
        }
        *cell = old_types[i];
    }
    free(old_types);
}

static struct var_def** var_cell(
    const struct scope* scope,
    const char* name
) {
    unsigned long long orig_idx = hash_name(name, scope->var_cap);
    unsigned long long idx = orig_idx;

    do {
        if (scope->vars[idx] == NULL
                || strcmp(name, scope->vars[idx]->name) == 0)
            return &scope->vars[idx];
    } while ((idx = (idx + 1) % scope->var_cap) != orig_idx);
    return NULL;
}

static void rehash_vars(struct scope* scope) {
    struct var_def** old_vars = scope->vars;
    unsigned long long old_cap = scope->var_cap;
    scope->var_cap *= 2;
    scope->vars = calloc(scope->var_cap, sizeof(struct var_def*));
    if (scope->vars == NULL) {
        fprintf(stderr, "ccc: out of memory\n");
        exit(1);
    }

    for (unsigned long long i = 0; i < old_cap; i++) {
        if (old_vars[i] == NULL) continue;

        struct var_def** cell = var_cell(scope, old_vars[i]->name);
        if (cell == NULL) {
            fprintf(stderr, "ccc: vars rehash failed, likely a bug\n");
            exit(1);
        }
        *cell = old_vars[i];
    }
    free(old_vars);
}

void scope_push(struct scope** p_scope) {
    struct scope* inner_scope = calloc(1, sizeof(struct scope));
    scope_init(inner_scope);
    inner_scope->next_out = *p_scope;
    *p_scope = inner_scope;
}

void scope_pop(struct scope** p_scope) {
    *p_scope = (*p_scope)->next_out;
}

void scope_install_default_types(struct ast* ast) {
    /* TODO: don't let people make void variables */
    ast->void_type = scope_define_type(ast->root_scope, (struct type_def) {
        .key = {
            .name = strdup("void"),
            .how_long = 0,
            .marked_signed = false,
            .marked_unsigned = false,
        },
        .is_signed = false,
        .is_floating = false,
        .sz = 0,
    });

    ast->char_type = scope_define_type(ast->root_scope, (struct type_def) {
        .key = {
            .name = strdup("char"),
            .how_long = 0,
            .marked_signed = false,
            .marked_unsigned = false,
        },
        .sz = 1,
        .is_signed = true, /* implementation defined babyyyyyy */
        .is_floating = false,
    });

    scope_define_type(ast->root_scope, (struct type_def) {
        .key = {
            .name = strdup("short"),
            .how_long = 0,
            .marked_signed = false,
            .marked_unsigned = false,
        },
        .sz = 2,
        .is_signed = true,
        .is_floating = false,
    });

    scope_define_type(ast->root_scope, (struct type_def) {
        .key = {
            .name = strdup("int"),
            .how_long = 0,
            .marked_signed = false,
            .marked_unsigned = false,
        },
        .sz = 4,
        .is_signed = true,
        .is_floating = false,
    });

    ast->integral_type = scope_define_type(ast->root_scope, (struct type_def) {
        .key = {
            .name = strdup("int"),
            .how_long = 1,
            .marked_signed = false,
            .marked_unsigned = false,
        },
        .sz = 8,
        .is_signed = true,
        .is_floating = false,
    });
    ast->string_type = ast->integral_type; /* TODO: support pointers */

    scope_define_type(ast->root_scope, (struct type_def) {
        .key = {
            .name = strdup("float"),
            .how_long = 0,
            .marked_signed = false,
            .marked_unsigned = false,
        },
        .sz = 4,
        .is_signed = true,
        .is_floating = true,
    });

    ast->decimal_type = scope_define_type(ast->root_scope, (struct type_def) {
        .key = {
            .name = strdup("double"),
            .how_long = 0,
            .marked_signed = false,
            .marked_unsigned = false,
        },
        .sz = 8,
        .is_signed = true,
        .is_floating = true,
    });
}

bool scope_get_type(
    const struct scope* scope,
    const struct type_def** p_entry,
    const struct type_key* key
) {
    for (; scope != NULL; scope = scope->next_out) {
        struct type_def** cell = type_cell(scope, key);
        if (cell == NULL || *cell == NULL) continue;
        if (p_entry != NULL) *p_entry = *cell;
        return true;
    }
    return false;
}

const struct type_def* scope_define_type(
    struct scope* scope,
    struct type_def type
) {
    struct type_def** cell = type_cell(scope, &type.key);
    while (cell == NULL) {
        rehash_types(scope);
        cell = type_cell(scope, &type.key);
    }

    /* redefinition leaks memory, so refuse */
    if (*cell != NULL) return NULL;

    *cell = calloc(1, sizeof(struct type_def));
    if (*cell == NULL) {
        fprintf(stderr, "ccc: out of memory\n");
        exit(1);
    }
    **cell = type;
    return *cell;
}

bool scope_get_var(
    const struct scope* scope,
    struct var_def** p_entry,
    const char* name
) {
    for (; scope != NULL; scope = scope->next_out) {
        struct var_def** cell = var_cell(scope, name);
        if (cell == NULL || *cell == NULL) continue;
        if (p_entry != NULL) *p_entry = *cell;
        return true;
    }
    return false;
}

struct var_def* scope_define_var(struct scope* scope, struct var_def var) {
    struct var_def** cell = var_cell(scope, var.name);
    while (cell == NULL) {
        rehash_vars(scope);
        cell = var_cell(scope, var.name);
    }

    /* redefinition leaks memory, so refuse */
    if (*cell != NULL) return NULL;

    if (*cell == NULL) {
        *cell = calloc(1, sizeof(struct var_def));
        if (*cell == NULL) {
            fprintf(stderr, "ccc: out of memory\n");
            exit(1);
        }
    }
    **cell = var;
    return *cell;
}