#include "type_checker.h"
#include "scope.h"
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <stddef.h>

#define TYPE_PANIC(format, ...) {\
    fprintf(\
        stderr,\
        "ccc: type error: " format "\n" __VA_OPT__(,)\
        __VA_ARGS__);\
    exit(1);\
}

static const struct ast* ast_ref;
static struct scope* scope;
static struct var_def* active_fn;

static void type_check_expr(struct expr_node* node);
static void type_check_stmt(struct stmt_node* node);
static void type_check_group(struct group_node* node);

static char* get_type_name(const struct type* type) {
    switch (type->type) {
        case TP_DATA:
            return strdup(type->data.data_type->name);
        case TP_PTR:
            char* stars = malloc((type->pointer.ptr_level + 1) * sizeof(char));
            memset(stars, '*', type->pointer.ptr_level);
            stars[type->pointer.ptr_level] = 0;
            char* name =
                ccc_sprintf("%s%s", type->pointer.data_type->name, stars);
            free(stars);
            return name;
        case TP_FN:
            /* TODO: un-stub :'( */
            return "function";
    }
    unreachable();
}

static bool is_callable(const struct type* type) {
    switch (type->type) {
        case TP_DATA:
        case TP_PTR:
            return false;
        case TP_FN:
            return true;
    }
    unreachable();
}

static void assert_cast_compatible(
    const struct type* lval_type,
    const struct type* rval_type
) {
    /* TODO: impl */
    /* TODO: we should also insert cast nodes eventually */
    if (false) {
        TYPE_PANIC(
            "cannot assign value of type '%s' to '%s'",
            get_type_name(lval_type),
            get_type_name(rval_type));
    }
}

static void enter_scope(struct scope* child_scope) {
    if (child_scope == NULL ||
            child_scope->next_out != scope)
        TYPE_PANIC("enter_scope: scopes are misaligned");

    scope = child_scope;
}

static void exit_scope(struct scope* child_scope) {
    if (child_scope != scope || child_scope->next_out == NULL)
        TYPE_PANIC("exit_scope: scopes are misaligned");

    scope = child_scope->next_out;
}

static const struct type* resolve_var_ref(struct var_ref_node* node) {
    if (node->def_ref->type == NULL)
        TYPE_PANIC("variable '%s' has undefined type.", node->def_ref->name);
    return node->def_ref->type;
}

static void type_check_decl(
    const struct type* decl_type,
    struct decl_node* node
) {
    node->def_ref->type = decl_type;

    if (node->initial_value != NULL) {
        type_check_expr(node->initial_value);
        assert_cast_compatible(decl_type, node->initial_value->resolved_type);
    }
}

static void type_check_decl_list(struct decl_list_node* node) {
    for (struct decl_node* cur = node->head; cur != NULL; cur = cur->next)
        type_check_decl(&node->type, cur);
}

static bool is_expr_assignable(struct expr_node* node) {
    switch (node->type) {
        case EXPR_VAR_REF:
            return true;
        case EXPR_PAREN:
            const struct expr_list_node* expr_list =
                node->inner.paren.expr_list;
            return expr_list->next == NULL
                && is_expr_assignable(expr_list->expr);
        default:
            return false;
    }
}

static const struct type* resolve_assign(struct assign_node* node) {
    if (!is_expr_assignable(node->lval))
        TYPE_PANIC("expression is not assignable");

    type_check_expr(node->lval);
    type_check_expr(node->rval);
    const struct type* lval_type = node->lval->resolved_type;
    const struct type* rval_type = node->rval->resolved_type;

    assert_cast_compatible(lval_type, rval_type);
    return lval_type;
}

static void type_check_expr_list(struct expr_list_node* node) {
    for (; node != NULL; node = node->next) {
        type_check_expr(node->expr);
        node->resolved_type = node->expr->resolved_type;
    }
}

static const struct type* resolve_call(struct call_node* node) {
    type_check_expr(node->expr);
    const struct fn_type* fn_type = &node->expr->resolved_type->fn;
    if (!is_callable(node->expr->resolved_type))
        TYPE_PANIC(
            "called object '%s' is not a function",
            get_type_name(node->expr->resolved_type));

    type_check_expr_list(node->args);
    struct type_list* arg_type = fn_type->arg_types;
    struct expr_list_node* arg_eval = node->args;
    while (arg_type != NULL && arg_eval != NULL) {
        const struct type* decl_type = arg_type->type;
        const struct type* eval_type = arg_eval->expr->resolved_type;
        assert_cast_compatible(decl_type, eval_type);

        arg_type = arg_type->next;
        arg_eval = arg_eval->next;
    }

    if (arg_type != NULL)
        TYPE_PANIC("missing arguments to function");
    if (arg_eval != NULL)
        TYPE_PANIC("too many arguments to function");

    return fn_type->return_type;
}

static const struct type* resolve_unary(struct unary_node* node) {
    type_check_expr(node->expr);
    /* TODO: delegate by individual operation to prohibit shit like -"yes" */
    return node->expr->resolved_type;
}

static const struct type* resolve_binary(struct binary_node* node) {
    /* TODO: math rules and such lol */
    type_check_expr(node->lhs);
    type_check_expr(node->rhs);
    assert_cast_compatible(node->lhs->resolved_type, node->rhs->resolved_type);
    return node->lhs->resolved_type;
}

static const struct type* resolve_paren(struct paren_node* node) {
    type_check_expr_list(node->expr_list);
    return node->expr_list->resolved_type;
}

static const struct type* resolve_cast(struct cast_node* node) {
    type_check_expr(node->expr);
    assert_cast_compatible(&node->type, node->expr->resolved_type);
    return &node->type;
}

static void type_check_expr(struct expr_node* node) {
    switch (node->type) {
        case EXPR_INT_LIT:
            node->resolved_type = &integral_type;
            break;
        case EXPR_FLOAT_LIT:
            node->resolved_type = &floating_type;
            break;
        case EXPR_CHAR_LIT:
            node->resolved_type = &character_type;
            break;
        case EXPR_STR_LIT:
            node->resolved_type = &string_type;
            break;
        case EXPR_VAR_REF:
            node->resolved_type = resolve_var_ref(&node->inner.var_ref);
            break;
        case EXPR_ASSIGN:
            node->resolved_type = resolve_assign(&node->inner.assign);
            break;
        case EXPR_CALL:
            node->resolved_type = resolve_call(&node->inner.call);
            break;
        case EXPR_UNARY:
            node->resolved_type = resolve_unary(&node->inner.unary);
            break;
        case EXPR_BINARY:
            node->resolved_type = resolve_binary(&node->inner.binary);
            break;
        case EXPR_PAREN:
            node->resolved_type = resolve_paren(&node->inner.paren);
            break;
        case EXPR_CAST:
            node->resolved_type = resolve_cast(&node->inner.cast);
            break;
    }
}

static void type_check_return(struct return_node* node) {
    if (active_fn == NULL) TYPE_PANIC("must be inside a function to return");
    const struct type* return_type = active_fn->type->fn.return_type;
    if (node->ret_val != NULL) {
        type_check_expr_list(node->ret_val);
        assert_cast_compatible(return_type, node->ret_val->resolved_type);
    } else if (return_type->type != TP_DATA
            || return_type->data.data_type != &void_type)
        TYPE_PANIC(
            "void function '%s' should not return a value",
            active_fn->name);
}

static void type_check_if(struct if_node* node) {
    enter_scope(node->scope);

    type_check_expr_list(node->cond);
    type_check_stmt(node->true_branch);
    if (node->false_branch != NULL) type_check_stmt(node->false_branch);

    exit_scope(node->scope);
}

static void type_check_loop_init(struct loop_init_node* node) {
    switch (node->type) {
        case INIT_EXPR_LIST:
            type_check_expr_list(node->expr_list);
            break;
        case INIT_DECL_LIST:
            type_check_decl_list(node->decl_list);
            break;
    }
}

static void type_check_loop(struct loop_node* node) {
    enter_scope(node->scope);

    if (node->init != NULL) type_check_loop_init(node->init);
    if (node->cond != NULL) type_check_expr_list(node->cond);
    if (node->incr != NULL) type_check_expr_list(node->incr);
    type_check_stmt(node->body);

    exit_scope(node->scope);
}

static void type_check_stmt(struct stmt_node* node) {
    switch (node->type) {
        case STMT_EMPTY:
            break;
        case STMT_EXPR_LIST:
            type_check_expr_list(&node->inner.expr_list);
            break;
        case STMT_DECL_LIST:
            type_check_decl_list(&node->inner.decl_list);
            break;
        case STMT_RETURN:
            type_check_return(&node->inner.return_);
            break;
        case STMT_GROUP:
            type_check_group(&node->inner.group);
            break;
        case STMT_IF:
            type_check_if(&node->inner.if_);
            break;
        case STMT_LOOP:
            type_check_loop(&node->inner.loop);
            break;
    }
}

static void type_check_group(struct group_node* node) {
    enter_scope(node->scope);

    struct stmt_node* stmt = node->head;
    while (stmt != NULL) {
        type_check_stmt(stmt);
        stmt = stmt->next;
    }

    exit_scope(node->scope);
}

/* TODO: type check redefinition arguments against each other */
/* TODO: wire the body to the definition either here or in the parser */
static void type_check_fn_decl(struct fn_decl_node* node) {
    if (!is_callable(&node->type))
        TYPE_PANIC("function definition does not have function type");
    if (active_fn != NULL)
        TYPE_PANIC(
            "can't define function '%s' inside function '%s'",
            node->name,
            active_fn->name);
    if (!scope_get_var(scope, &active_fn, node->name))
        TYPE_PANIC("scopes are borked, missing symbol: '%s'", node->name);

    assert_cast_compatible(active_fn->type, &node->type);

    enter_scope(node->scope);
    if (node->body != NULL) {
        type_check_group(node->body);
        active_fn->fn.resolved = true;
    }
    exit_scope(node->scope);

    active_fn = NULL;
}

static void type_check_root(struct root_node* node) {
    switch (node->type) {
        case ROOT_FN_DECL:
            type_check_fn_decl(&node->inner.fn_decl);
            break;
    }
}

void type_check(struct ast* ast) {
    ast_ref = ast;
    scope = ast->root_scope;

    struct root_node* node = ast->root_node;
    while (node != NULL) {
        type_check_root(node);
        node = node->next;
    }

    ast_ref = NULL;
    scope = NULL;
}