#include <scc_ast_utils.h>
#include <scc_pos_log.h>
#include <scc_sema.h>
#include <sema_symtab.h>

static scc_ast_qual_type_t *got_type_callback(scc_sema_ctx_t *sema_ctx,
                                              const char *name);

static void gen_symbol_name(const scc_ast_decl_t *decl, scc_str_t *name) {
    switch (decl->base.type) {
    case SCC_AST_DECL_ENUM:
        *name = scc_str_from_cstr("$E_");
        scc_str_append_cstr(name, decl->name, scc_strlen(decl->name));
        break;
    case SCC_AST_DECL_STRUCT:
        *name = scc_str_from_cstr("$S_");
        scc_str_append_cstr(name, decl->name, scc_strlen(decl->name));
        break;
    case SCC_AST_DECL_UNION:
        *name = scc_str_from_cstr("$U_");
        scc_str_append_cstr(name, decl->name, scc_strlen(decl->name));
        break;
    default:
        break;
    }
}

static void symtab_add_symbol(scc_sema_symtab_t *sema_symtab,
                              scc_ast_decl_t *decl,
                              scc_ast_node_t *ast_node_ref) {

    if (decl->name == nullptr)
        return;
    // FIXME memory leak
    scc_str_t name = scc_str_empty();
    gen_symbol_name(decl, &name);
    if (scc_str_is_empty(&name)) {
        scc_sema_symtab_add_symbol(sema_symtab, decl->name, ast_node_ref);
    } else {
        scc_sema_symtab_add_symbol(sema_symtab, scc_str_as_cstr(&name),
                                   ast_node_ref);
    }
}

static void type_callback(scc_sema_ctx_t *sema_ctx,
                          scc_ast_node_kind_t node_type, void *node) {
    scc_sema_symtab_t *sema_symtab = sema_ctx->context;
    scc_ast_qual_type_t *type = SCC_AST_CAST_TO(scc_ast_qual_type_t, node);

    if (node_type == SCC_AST_TYPE_STRUCT) {
        symtab_add_symbol(sema_symtab, type->type->record.decl, &type->base);
    } else if (node_type == SCC_AST_TYPE_UNION) {
        symtab_add_symbol(sema_symtab, type->type->record.decl, &type->base);
    } else if (node_type == SCC_AST_TYPE_ENUM) {
        symtab_add_symbol(sema_symtab, type->type->record.decl, &type->base);
    }
    return;
}

static void expr_callback(scc_sema_ctx_t *sema_ctx,
                          scc_ast_node_kind_t node_type, void *node) {
    scc_sema_symtab_t *sema_symtab = sema_ctx->context;

    if (node_type == SCC_AST_UNKNOWN || node == nullptr) {
        return;
    }
    scc_ast_expr_t *expr = SCC_AST_CAST_TO(scc_ast_expr_t, node);
    if (node_type == SCC_AST_EXPR_IDENTIFIER) {
        scc_ast_node_t *node =
            scc_sema_symtab_lookup_symbol(sema_symtab, expr->identifier.name);
        if (node == nullptr) {
            SCC_ERROR(expr->base.loc, "sema error: Identifier '%s' not found",
                      expr->identifier.name);
        } else if (!SCC_AST_IS_A(scc_ast_decl_t, node)) {
            SCC_ERROR(expr->base.loc,
                      "sema error: Identifier '%s' is not a variable",
                      expr->identifier.name);
        } else {
            expr->identifier._target = SCC_AST_CAST_TO(scc_ast_decl_t, node);
        }
    }

    return;
}

static void stmt_callback(scc_sema_ctx_t *sema_ctx,
                          scc_ast_node_kind_t node_type, void *node) {
    scc_sema_symtab_t *sema_symtab = sema_ctx->context;

    if (node_type == scc_ast_stmt_t_BEGIN) {
        if (node == nullptr) {
            scc_sema_symtab_enter_scope(sema_symtab);
            return;
        }
        switch (((scc_ast_stmt_t *)node)->base.type) {
        case SCC_AST_STMT_WHILE:
        case SCC_AST_STMT_DO_WHILE:
        case SCC_AST_STMT_FOR:
            scc_vec_push(sema_ctx->break_stack, (scc_ast_stmt_t *)node);
            scc_vec_push(sema_ctx->continue_stack, (scc_ast_stmt_t *)node);
            break;
        case SCC_AST_STMT_SWITCH:
            scc_vec_push(sema_ctx->break_stack, (scc_ast_stmt_t *)node);
            break;
        default:
            Panic("Unhandled statement type");
            break;
        }
        return;
    } else if (node_type == scc_ast_stmt_t_END) {
        if (node == nullptr) {
            scc_sema_symtab_leave_scope(sema_symtab);
            return;
        }
        switch (((scc_ast_stmt_t *)node)->base.type) {
        case SCC_AST_STMT_WHILE:
        case SCC_AST_STMT_DO_WHILE:
        case SCC_AST_STMT_FOR:
            scc_vec_pop(sema_ctx->break_stack);
            scc_vec_pop(sema_ctx->continue_stack);
            break;
        case SCC_AST_STMT_SWITCH:
            scc_vec_pop(sema_ctx->break_stack);
            break;
        default:
            Panic("Unhandled statement type");
            break;
        }
        return;
    }
    if (node_type == SCC_AST_UNKNOWN || node == nullptr) {
        return;
    }
    scc_ast_stmt_t *stmt = SCC_AST_CAST_TO(scc_ast_stmt_t, node);
    switch (stmt->base.type) {
    case SCC_AST_STMT_BREAK: {
        if (scc_vec_size(sema_ctx->break_stack) == 0)
            SCC_ERROR(stmt->base.loc, "break not in loop/switch");
        else
            stmt->jump._target = scc_vec_at(
                sema_ctx->break_stack, scc_vec_size(sema_ctx->break_stack) - 1);
    } break;
    case SCC_AST_STMT_CONTINUE: {
        if (scc_vec_size(sema_ctx->continue_stack) == 0)
            SCC_ERROR(stmt->base.loc, "continue not in loop");
        else
            stmt->jump._target =
                scc_vec_at(sema_ctx->continue_stack,
                           scc_vec_size(sema_ctx->continue_stack) - 1);
    } break;
    case SCC_AST_STMT_GOTO: {
        scc_ast_stmt_t *goto_stmt = SCC_AST_CAST_TO(scc_ast_stmt_t, node);
        scc_sema_label_scope_t *scope = &sema_ctx->label_scope;
        if (scope->function_scope_depth != 1) {
            SCC_ERROR(goto_stmt->base.loc, "goto not inside any function");
            break;
        }

        // 尝试直接查找标签
        scc_ast_node_t *target =
            scc_hashtable_get(&scope->labels, goto_stmt->goto_stmt.label);
        if (target) {
            goto_stmt->goto_stmt._target =
                SCC_AST_CAST_TO(scc_ast_stmt_t, target);
        } else {
            // 未找到则加入待处理列表
            scc_vec_push(scope->pending_gotos, goto_stmt);
        }
    } break;
    case SCC_AST_STMT_LABEL: {
        scc_ast_stmt_t *label_stmt = SCC_AST_CAST_TO(scc_ast_stmt_t, node);
        scc_sema_label_scope_t *scope = &sema_ctx->label_scope;
        // 获取当前函数标签作用域
        if (scope->function_scope_depth != 1) {
            SCC_ERROR(label_stmt->base.loc, "label not inside any function");
            break;
        }

        // 检查重复定义
        if (scc_hashtable_get(&scope->labels, label_stmt->label_stmt.label)) {
            SCC_ERROR(label_stmt->base.loc, "duplicate label '%s'",
                      label_stmt->label_stmt.label);
            break;
        }

        // 添加标签
        scc_hashtable_set(&scope->labels, label_stmt->label_stmt.label,
                          &label_stmt->base);

        // 检查是否有等待该标签的 goto
        scc_vec_foreach(scope->pending_gotos, i) {
            scc_ast_stmt_t *goto_stmt = scc_vec_at(scope->pending_gotos, i);
            if (goto_stmt == nullptr)
                continue;
            if (scc_strcmp(goto_stmt->goto_stmt.label,
                           label_stmt->label_stmt.label) == 0) {
                goto_stmt->goto_stmt._target = label_stmt;
                scc_vec_at(scope->pending_gotos, i) = nullptr;
            }
        }
    } break;
    default:
        break;
    }
    return;
}

static void decl_callback(scc_sema_ctx_t *sema_ctx,
                          scc_ast_node_kind_t node_type, void *node) {
    scc_sema_symtab_t *sema_symtab = sema_ctx->context;

    // Function declaration scope
    // FIXME 使用 node_type 区分其他未来的scope
    if (node_type == scc_ast_decl_t_BEGIN) {
        // 创建标签作用域
        scc_sema_label_scope_t *scope = &sema_ctx->label_scope;
        scope->function_scope_depth += 1;
        scc_hashtable_cstr_init(&scope->labels);
        scc_vec_init(scope->pending_gotos);

        // 进入普通符号表作用域（用于参数和局部变量）
        scc_sema_symtab_enter_scope(sema_symtab);
        return;
    } else if (node_type == scc_ast_decl_t_END) {
        scc_sema_label_scope_t *scope = &sema_ctx->label_scope;
        scope->function_scope_depth -= 1;

        // 处理尚未绑定的 goto
        scc_vec_foreach(scope->pending_gotos, i) {
            scc_ast_stmt_t *goto_stmt = scc_vec_at(scope->pending_gotos, i);
            if (goto_stmt == nullptr) {
                continue;
            }
            // 标签未定义 -> 报错
            SCC_ERROR(goto_stmt->base.loc, "label '%s' used but not defined",
                      goto_stmt->goto_stmt.label);
        }

        // 清理
        scc_hashtable_drop(&scope->labels);
        scc_vec_free(scope->pending_gotos);

        // 退出普通符号表作用域
        scc_sema_symtab_leave_scope(sema_symtab);
        return;
    }
    if (node_type == SCC_AST_UNKNOWN || node == nullptr) {
        return;
    }
    scc_ast_decl_t *decl = SCC_AST_CAST_TO(scc_ast_decl_t, node);

    if (node_type == SCC_AST_DECL_TYPEDEF) {
        if (decl->name == nullptr) {
            SCC_ERROR(decl->base.loc, "typedef without name");
            return;
        }
        symtab_add_symbol(sema_symtab, decl, &decl->typedef_decl.type->base);
    } else if (node_type == SCC_AST_DECL_VAR) {
        symtab_add_symbol(sema_symtab, decl, &decl->base);
    } else if (node_type == SCC_AST_DECL_PARAM) {
        if (decl->name == nullptr) {
            if (decl->param.type->base.type == SCC_AST_TYPE_BUILTIN &&
                (decl->param.type->type->builtin.type ==
                     SCC_AST_BUILTIN_TYPE_VA_LIST ||
                 decl->param.type->type->builtin.type ==
                     SCC_AST_BUILTIN_TYPE_VOID)) {
                return;
            }
            SCC_ERROR(decl->base.loc, "sema error: Parameter must have a name");
            return;
        }
        symtab_add_symbol(sema_symtab, decl, &decl->base);
    } else if (node_type == SCC_AST_DECL_FUNC) {
        if (decl->name == nullptr) {
            SCC_ERROR(decl->base.loc, "sema error: Function must have a name");
        } else {
            // FIXME 重名函数...
            symtab_add_symbol(sema_symtab, decl, &decl->base);
        }
    }
    return;
}

static scc_ast_qual_type_t *got_type_callback(scc_sema_ctx_t *sema_ctx,
                                              const char *name) {
    scc_sema_symtab_t *sema_symtab = sema_ctx->context;
    scc_ast_node_t *node = scc_sema_symtab_lookup_symbol(sema_symtab, name);
    if (SCC_AST_IS_A(scc_ast_qual_type_t, node)) {
        return (scc_ast_qual_type_t *)node;
    }
    return nullptr;
}

void scc_sema_init(scc_sema_ctx_t *sema_ctx, scc_ast_ctx_t *ast_ctx) {
    sema_ctx->ast_ctx = ast_ctx;
    scc_sema_symtab_t *sema_symtab = scc_malloc(sizeof(scc_sema_symtab_t));
    if (sema_symtab == nullptr) {
        LOG_FATAL("out of memory");
        return;
    }
    sema_ctx->context = sema_symtab;
    sema_ctx->on_decl = decl_callback;
    sema_ctx->on_expr = expr_callback;
    sema_ctx->on_stmt = stmt_callback;
    sema_ctx->on_type = type_callback;
    sema_ctx->got_type = got_type_callback;
    sema_ctx->label_scope = (scc_sema_label_scope_t){0};

    scc_sema_symtab_init(sema_symtab);

    scc_vec_init(sema_ctx->break_stack);
    scc_vec_init(sema_ctx->continue_stack);

    // FIXME memory leak

    scc_ast_qual_type_t *type = SCC_AST_ALLOC_QUAL_TYPE(sema_ctx->ast_ctx);
    scc_ast_type_builtin_init(type, sema_ctx->ast_ctx,
                              SCC_AST_BUILTIN_TYPE_VA_LIST, scc_pos_create());
    scc_sema_symtab_add_symbol(sema_symtab, "__scc_builtin_va_list",
                               &type->base);

    scc_ast_decl_t *decl = SCC_AST_ALLOC_DECL(sema_ctx->ast_ctx);
    scc_ast_decl_val_init(decl, type, "__scc_builtin__", nullptr,
                          scc_pos_create());
    scc_sema_symtab_add_symbol(sema_symtab, "__func__", &decl->base);

    scc_ast_qual_type_t *built_func_type =
        SCC_AST_ALLOC_QUAL_TYPE(sema_ctx->ast_ctx);
    // FIXME hack
    built_func_type->base.type = SCC_AST_TYPE_FUNCTION;

    scc_ast_decl_t *builin_func = SCC_AST_ALLOC_DECL(sema_ctx->ast_ctx);
    scc_ast_decl_func_init(builin_func, built_func_type, "__scc_builtin_func",
                           nullptr, scc_pos_create());
    scc_sema_symtab_add_symbol(sema_symtab, "__scc_builtin_va_start",
                               &builin_func->base);
    scc_sema_symtab_add_symbol(sema_symtab, "__scc_builtin_va_end",
                               &builin_func->base);
    scc_sema_symtab_add_symbol(sema_symtab, "__scc_builtin_va_arg",
                               &builin_func->base);
    scc_sema_symtab_add_symbol(sema_symtab, "__scc_builtin_va_copy",
                               &builin_func->base);
}

void scc_sema_drop(scc_sema_ctx_t *sema_ctx) {
    Assert(scc_vec_size(sema_ctx->break_stack) == 0 &&
           scc_vec_size(sema_ctx->continue_stack) == 0);
    scc_vec_free(sema_ctx->break_stack);
    scc_vec_free(sema_ctx->continue_stack);

    // FIXME drop obj
    if (sema_ctx->context) {
        scc_sema_symtab_drop(sema_ctx->context);
    }
}