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

static void type_spec_print(FILE* file, TypeSpec spec)
{

    struct {
        TypeSpec spec;
        const char* ident;
    } spec_idents[] = {
        { TypeSpec_Const, "const" },
        { TypeSpec_Short, "short" },
        { TypeSpec_Long, "long" },
        { TypeSpec_Long2, "long" },
        { TypeSpec_Unsigned, "unsigned" },
    };

    size_t spec_idents_size = sizeof(spec_idents) / sizeof(spec_idents[0]);

    for (size_t i = 0; i < spec_idents_size; ++i) {
        if (spec & spec_idents[i].spec) {
            fprintf(file, "%s ", spec_idents[i].ident);
        }
    }
}

void type_print(FILE* file, const Type* type)
{
    type_spec_print(file, type->spec);

    struct {
        TypeTy ty;
        const char* ident;
    } ty_idents[] = {
        { Type_Int, "int" },
        { Type_Char, "char" },
        { Type_Float, "float" },
        { Type_Double, "double" },
        { Type_Bool, "bool" },
        { Type_Void, "void" },
    };

    size_t ty_idents_size = sizeof(ty_idents) / sizeof(ty_idents[0]);

    for (size_t i = 0; i < ty_idents_size; ++i) {
        if (type->ty == ty_idents[i].ty) {
            fprintf(file, "%s", ty_idents[i].ident);
            return;
        }
    }

    if (type->ty == Type_Pointer) {
        fprintf(file, " * ");
        type_spec_print(file, type->pointer.spec);
        type_print(file, type->pointer.inner);
        return;
    }

    assert(false && "not implemented");
}

typedef struct {
    Hashmap ord_syms;
    Hashmap tag_syms;
} Rib;

static void rib_init(Rib* r)
{
    *r = (Rib) {
        .ord_syms = { },
        .tag_syms = { },
    };
    hashmap_init(&r->ord_syms, sizeof(Sym*));
    hashmap_init(&r->tag_syms, sizeof(Sym*));
}

static void rib_deinit(Rib* r)
{
    hashmap_deinit(&r->ord_syms);
    hashmap_deinit(&r->tag_syms);
}

struct Sema {
    Arena sym_arena;
    Arena type_arena;
    ARRAY_FIELD(Rib, ribs);
    Hashmap labels;
    Keymap expr_types;
};

Sema* sema_new(void)
{
    Sema* s = malloc(sizeof(Sema));
    *s = (Sema) {
        .sym_arena = { },
        .type_arena = { },
        ARRAY_FIELD_INIT(ribs),
        .labels = { },
        .expr_types = { },
    };

    arena_init(&s->sym_arena, 1024);
    arena_init(&s->type_arena, 1024);

    ARRAY_FIELD_PUSH(s, ribs, (&(Rib) { }));
    rib_init(&s->ribs[0]);

    hashmap_init(&s->labels, sizeof(LabelSym*));
    keymap_init(&s->expr_types, sizeof(Type*));

    return s;
}

void sema_free(Sema* s)
{
    arena_deinit(&s->sym_arena);
    arena_deinit(&s->type_arena);
    for (size_t i = 0; i < s->ribs_size; ++i)
        rib_deinit(&s->ribs[i]);
    if (s->ribs)
        free(s->ribs);
    hashmap_deinit(&s->labels);
    keymap_deinit(&s->expr_types);
    free(s);
}

void sema_scope_enter(Sema* s)
{
    ARRAY_FIELD_PUSH(s, ribs, (&(Rib) { }));
    rib_init(&s->ribs[s->ribs_size - 1]);
}

void sema_scope_leave(Sema* s)
{
    assert(s->ribs_size > 0);
    rib_deinit(&s->ribs[s->ribs_size - 1]);
    s->ribs_size -= 1;
}

static Sym* get_ord_sym(Sema* s, const char* ident)
{
    Sym* sym = nullptr;
    for (size_t i = 0; i < s->ribs_size; ++i)
        if ((sym = hashmap_get(&s->ribs[s->ribs_size - i - 1].ord_syms, ident)))
            break;
    return sym;
}

static void error_if_multi_spec(
    bool type_specced, Loc loc, const Type* first, TypeTy second_ty)
{
    if (type_specced) {
        fprintf(stderr, "error: cannot combine type '");
        type_print(stderr, &(Type) { .ty = second_ty });
        fprintf(stderr, "' with type '");
        type_print(stderr, first);
        fprintf(stderr, "' on line %d\n", loc.line);
        exit(1);
    }
}

static void decode_decl_spec(Type* type, AstSpec as, Loc loc)
{
    bool type_specced = false;

    if (as & AstSpec_Int) {
        error_if_multi_spec(type_specced, loc, type, Type_Char);
        type_specced = true;

        type->ty = Type_Int;
    }
    if (as & AstSpec_Char) {
        error_if_multi_spec(type_specced, loc, type, Type_Char);
        type_specced = true;

        type->ty = Type_Char;
    }
    if (as & AstSpec_Void) {
        error_if_multi_spec(type_specced, loc, type, Type_Void);
        type_specced = true;

        type->ty = Type_Void;
    }
    if (as & (AstSpec_Enum | AstSpec_Struct | AstSpec_Union)) {
        assert(false && "not implemented");
    }

    if (as & AstSpec_Unsigned) {
        type->spec |= TypeSpec_Unsigned;
    }
    if (as & AstSpec_Const) {
        type->spec |= TypeSpec_Const;
    }
}

static Type* check_decl_type(Sema* s, AstDecl* decl)
{
    Type* type = arena_alloc(&s->type_arena, sizeof(Type));
    *type = (Type) { };

    decode_decl_spec(type, decl->spec, decl->loc);

    return type;
}

static void error_if_redefinition(Hashmap* map, const char* ident, Loc loc)
{
    Sym* sym;
    if ((sym = hashmap_get(map, ident))) {
        fprintf(stderr,
            "error: redefinition of '%s' on line %d\n",
            ident,
            loc.line);
        exit(1);
    }
}

static void check_declar(Sema* s,
    AstDeclar* declar,
    Type* type,
    Type** out_declar_type,
    const char** out_declar_ident)
{
    switch (declar->ty) {
        case AstDeclar_Ident: {
            *out_declar_type = type;
            *out_declar_ident = declar->ident;
            return;
        }
        case AstDeclar_Pointer: {
            Type* new_type = arena_alloc(&s->type_arena, sizeof(Type));
            *new_type = (Type) {
                .ty = Type_Pointer,
                .spec = 0,
                .pointer = {
                    .inner = type,
                    .spec = 0,
                },
            };
            decode_decl_spec(new_type, declar->pointer.spec, declar->loc);
            check_declar(s,
                declar->pointer.declar,
                new_type,
                out_declar_type,
                out_declar_ident);
            return;
        }
        case AstDeclar_Array: {
            assert(false && "not implemented");
            break;
        }
        case AstDeclar_Function: {
            Type* new_type = arena_alloc(&s->type_arena, sizeof(Type));
            *new_type = (Type) {
                .ty = Type_Function,
                .spec = 0,
                .function = {
                    .return_type = type,
                    ARRAY_FIELD_INIT(params),
                },
            };

            for (size_t i = 0; i < declar->function.params_size; ++i) {
                AstDecl* param_decl = declar->function.params[i];

                assert(param_decl->declars_size == 1);
                AstDeclar* param_declar = param_decl->declars[0];

                Type* decl_type = check_decl_type(s, param_decl);

                const char* declar_ident;
                Type* declar_type;
                check_declar(
                    s, param_declar, decl_type, &declar_type, &declar_ident);

                (void)declar_ident;

                ARENA_ARRAY_FIELD_PUSH(
                    &s->type_arena, &new_type->function, params, &declar_type);
            }

            // TODO: handle forward declarations

            check_declar(s,
                declar->pointer.declar,
                new_type,
                out_declar_type,
                out_declar_ident);
            return;
        }
    }
    assert(false && "unreachable");
}

void sema_function_enter(Sema* s, AstSpec spec, AstDeclar* declar)
{
    assert(declar->ty == AstDeclar_Function);

    // clear label map
    hashmap_deinit(&s->labels);
    hashmap_init(&s->labels, sizeof(LabelSym*));

    Type* type = arena_alloc(&s->type_arena, sizeof(Type));
    *type = (Type) { };
    decode_decl_spec(type, spec, declar->loc);

    Rib* rib = &s->ribs[s->ribs_size - 1];
    error_if_redefinition(&rib->ord_syms, declar->ident, declar->loc);

    Sym* sym = arena_alloc(&s->sym_arena, sizeof(Sym));
    *sym = (Sym) {
        .ty = Sym_Declar,
        .declar = {
            .type = type,
        },
    };
    hashmap_set(&rib->ord_syms, declar->ident, &sym);

    sema_scope_enter(s);

    for (size_t i = 0; i < declar->function.params_size; ++i) {
        AstDecl* param_decl = declar->function.params[i];

        assert(param_decl->declars_size == 1);
        AstDeclar* param_declar = param_decl->declars[0];

        Type* decl_type = check_decl_type(s, param_decl);

        const char* declar_ident;
        Type* declar_type;
        check_declar(s, param_declar, decl_type, &declar_type, &declar_ident);

        Rib* rib = &s->ribs[s->ribs_size - 1];
        error_if_redefinition(&rib->ord_syms, declar->ident, declar->loc);

        Sym* sym = arena_alloc(&s->sym_arena, sizeof(Sym));
        *sym = (Sym) {
            .ty = Sym_Param,
            .param = {
                .type = declar_type,
                .idx = i,
            },
        };
        hashmap_set(&rib->ord_syms, declar_ident, &sym);
    }
}

void sema_function_leave(Sema* s)
{
    sema_scope_leave(s);
}

void sema_visit_stmt(Sema* s, AstStmt* stmt)
{
    if (stmt->ty == AstStmt_Decl) {
        Type* type = check_decl_type(s, stmt->decl);

        if (stmt->decl->function_body) {
            // assume already handled
            return;
        }

        for (size_t i = 0; i < stmt->decl->declars_size; ++i) {
            AstDeclar* declar = stmt->decl->declars[i];

            Type* declar_type;
            const char* declar_ident;
            check_declar(s, declar, type, &declar_type, &declar_ident);

            Rib* rib = &s->ribs[s->ribs_size - 1];
            error_if_redefinition(&rib->ord_syms, declar->ident, declar->loc);

            Sym* sym = arena_alloc(&s->sym_arena, sizeof(Sym));
            *sym = (Sym) {
                .ty = Sym_Declar,
                .declar = {
                    .type = declar_type,
                },
            };
            hashmap_set(&rib->ord_syms, declar_ident, &sym);
        }
    }
}

void sema_visit_expr(Sema* s, AstExpr* expr)
{
    Type* type = nullptr;

    switch (expr->ty) {
        case AstExpr_Ident: {
            Sym* sym;
            if ((sym = get_ord_sym(s, expr->ident)) == nullptr) {
                fprintf(stderr,
                    "error: undefined identifier '%s' on line %d\n",
                    expr->ident,
                    expr->loc.line);
                exit(1);
            }

            switch (sym->ty) {
                case Sym_Declar:
                    type = sym->declar.type;
                    break;
                case Sym_Param:
                    type = sym->param.type;
                    break;
            }

            break;
        }
        case AstExpr_Str: {
            Type* inner_type = arena_alloc(&s->type_arena, sizeof(Type));
            *inner_type = (Type) {
                .ty = Type_Char,
                .spec = TypeSpec_Const,
            };

            type = arena_alloc(&s->type_arena, sizeof(Type));
            *type = (Type) {
                .ty = Type_Pointer,
                .pointer = {
                    .inner = inner_type,
                    .spec = 0,
                },
            };
            break;
        };
        case AstExpr_Call: {
            Type* inner_type
                = keymap_get(&s->expr_types, (size_t)expr->call.expr);
            assert(inner_type);
            printf("inner_type->ty == %d\n", inner_type->ty);
            assert(inner_type->ty == Type_Function);
            break;
        }
    }

    keymap_set(&s->expr_types, (KeymapKey)expr, &type);
}