#include <ir_builtin.h>
#include <scc_ast2ir.h>

static scc_ir_node_ref_t ast_expr_to_ir(scc_ir_builder_t *ctx,
                                        scc_ast_expr_t *expr);
static void ast_stmt_to_ir(scc_ir_builder_t *ctx, scc_ast_stmt_t *stmt);
static void ast_decl_to_ir(scc_ir_builder_t *ctx, scc_ast_decl_t *decl);
static scc_ir_type_ref_t ast_type_to_ir_type(scc_ir_builder_t *ctx,
                                             scc_ast_type_t *ast_type);

// 转换AST类型为IR类型
static scc_ir_type_ref_t ast_type_to_ir_type(scc_ir_builder_t *ctx,
                                             scc_ast_type_t *ast_type) {
    if (ctx == null || ast_type == null) {
        LOG_ERROR("args is null");
        return 0;
    }
    scc_ir_type_t ir_type;

    switch (ast_type->base.type) {
    case SCC_AST_TYPE_BUILTIN: {
        // 映射内置类型
        scc_ir_type_init(&ir_type, SCC_IR_TYPE_I32);
        // TODO: 根据具体内置类型设置
        break;
    }

    case SCC_AST_TYPE_POINTER: {
        scc_ir_type_init(&ir_type, SCC_IR_TYPE_PTR);

        // 创建指向类型并添加到程序类型列表
        scc_ir_type_ref_t pointee_type =
            ast_type_to_ir_type(ctx, ast_type->pointer.pointee);

        // 注意：我们需要找到一种合适的方式来存储类型信息
        // 目前的IR设计中类型信息应该直接存储在类型结构中
        ir_type.data.pointer.base = pointee_type;
        break;
    }

    case SCC_AST_TYPE_ARRAY: {
        scc_ir_type_init(&ir_type, SCC_IR_TYPE_ARRAY);

        // 创建元素类型并添加到程序类型列表
        scc_ir_type_ref_t element_type =
            ast_type_to_ir_type(ctx, ast_type->array.element);

        // 将类型添加到程序的类型容器中
        ir_type.data.array.base = element_type;
        // TODO: 处理数组大小表达式
        ir_type.data.array.len = 0; // 暂时设为0
        break;
    }

    case SCC_AST_TYPE_FUNCTION: {
        scc_ir_type_init(&ir_type, SCC_IR_TYPE_FUNC);

        // 处理返回类型
        scc_ir_type_ref_t ret_type =
            ast_type_to_ir_type(ctx, ast_type->function.return_type);

        // 将返回类型添加到程序的类型容器中
        ir_type.data.function.ret_type = ret_type;

        // 转换参数类型
        scc_ir_type_ref_vec_t params;
        scc_vec_init(params);
        scc_vec_foreach(ast_type->function.param_types, i) {
            scc_ast_type_t *param_type =
                scc_vec_at(ast_type->function.param_types, i);
            scc_ir_type_ref_t tmp_type = ast_type_to_ir_type(ctx, param_type);
            scc_vec_push(params, tmp_type);
        }
        ir_type.data.function.params = params;
        break;
    }

    default:
        LOG_FATAL("Unsupported AST type: %d", ast_type->base.type);
        return 0;
    }
    return scc_ir_ctx_new_type(&ctx->ctx, &ir_type);
}

// 转换AST表达式为IR节点
static scc_ir_node_ref_t ast_expr_to_ir(scc_ir_builder_t *ctx,
                                        scc_ast_expr_t *expr) {
    if (ctx == null || expr == null) {
        LOG_ERROR("args is null");
        return 0;
    }

    switch (expr->base.type) {
    case SCC_AST_EXPR_IDENTIFIER: {
        // TODO maybe error or need symtab
        scc_ir_node_t in;
        scc_ir_node_init(&in, expr->identifier.name, SCC_IR_NODE_LOAD);
        return scc_ir_ctx_new_node(&ctx->ctx, &in);
    }

    case SCC_AST_EXPR_INT_LITERAL: {
        return scc_ir_ctx_get_i32_const(&ctx->ctx, expr->literal.value.i);
    }

    case SCC_AST_EXPR_BINARY: {
        // 转换左右操作数
        scc_ir_node_ref_t lhs, rhs;
        lhs = ast_expr_to_ir(ctx, expr->binary.lhs);
        rhs = ast_expr_to_ir(ctx, expr->binary.rhs);

        // 映射操作符
        scc_ir_op_type_t op;
        switch (expr->binary.op) {
            /* clang-format off */
        case SCC_AST_OP_ADD: op = IR_OP_ADD; break;
        case SCC_AST_OP_SUB: op = IR_OP_SUB; break;
        case SCC_AST_OP_MUL: op = IR_OP_MUL; break;
        case SCC_AST_OP_DIV: op = IR_OP_DIV; break;
        case SCC_AST_OP_MOD: op = IR_OP_MOD; break;
        case SCC_AST_OP_LOGICAL_AND: op = IR_OP_AND; break;
        case SCC_AST_OP_LOGICAL_OR:  op = IR_OP_OR;  break;
        case SCC_AST_OP_BITWISE_XOR: op = IR_OP_XOR; break;
        case SCC_AST_OP_LEFT_SHIFT:  op = IR_OP_SHL; break;
        case SCC_AST_OP_RIGHT_SHIFT: op = IR_OP_SHR; break;
        case SCC_AST_OP_EQUAL:       op = IR_OP_EQ;  break;
        case SCC_AST_OP_NOT_EQUAL:   op = IR_OP_NEQ; break;
        case SCC_AST_OP_LESS:        op = IR_OP_LT;  break;
        case SCC_AST_OP_LESS_EQUAL:  op = IR_OP_LE;  break;
        case SCC_AST_OP_GREATER:     op = IR_OP_GT;  break;
        case SCC_AST_OP_GREATER_EQUAL: op = IR_OP_GE;  break;
        case SCC_AST_OP_ASSIGN: {
            // 赋值表达式：存储右值到左值位置
            return scc_ir_builder_store(ctx, lhs, rhs);
            
        }
        case SCC_AST_OP_ASSIGN_ADD:
            TODO();
            /* clang-format on */
        default:
            LOG_WARN("Unsupported binary operator: %d", expr->binary.op);
            op = IR_OP_ADD; // 默认
        }

        // 创建操作节点
        return scc_ir_builder_binop(ctx, op, lhs, rhs);
    }

    case SCC_AST_EXPR_UNARY: {
        // 转换操作数
        scc_ir_node_ref_t operand = ast_expr_to_ir(ctx, expr->unary.operand);

        // 映射一元操作符
        switch (expr->unary.op) {
        case SCC_TOK_SUB:
            // 负号
            // 实现为0 - operand
            return scc_ir_builder_binop(ctx, IR_OP_SUB,
                                        scc_ir_ctx_get_builtin_zero(&ctx->ctx),
                                        operand);
        case SCC_TOK_BIT_NOT:
            // 按位取反
            return scc_ir_builder_binop(ctx, IR_OP_NOT, operand, 0);
        case SCC_TOK_NOT:
            // 逻辑非
            // 实现为与0比较
            return scc_ir_builder_binop(
                ctx, IR_OP_EQ, scc_ir_ctx_get_builtin_zero(&ctx->ctx), operand);
        default:
            LOG_WARN("Unsupported unary operator: %d", expr->unary.op);
            return 0;
        }
    }

    case SCC_AST_EXPR_COND: {
        TODO();
        break;
    }

    case SCC_AST_EXPR_CALL: {
        // 转换参数
        scc_ir_node_ref_vec_t args;
        scc_vec_init(args);

        // 检查参数是否为空
        if (expr->call.args.data != null) {
            scc_vec_foreach(expr->call.args, i) {
                scc_ast_expr_t *arg_expr = scc_vec_at(expr->call.args, i);
                scc_ir_node_ref_t arg_node;
                arg_node = ast_expr_to_ir(ctx, arg_expr);
                scc_vec_push(args, arg_node);
            }
        }

        // 创建调用节点（需要查找函数定义）
        // TODO: 需要符号表查找函数
        scc_ir_node_ref_t func =
            scc_ir_builder_call(ctx, 0, args.data, args.size);
        scc_vec_free(args);
        return func;
    }

    default:
        LOG_WARN("Unsupported expression type: %d", expr->base.type);
        return 0;
    }
}

// 转换AST语句为IR
static void ast_stmt_to_ir(scc_ir_builder_t *ctx, scc_ast_stmt_t *stmt) {
    if (stmt == null) {
        return;
    }

    switch (stmt->base.type) {
    case SCC_AST_STMT_COMPOUND: {
        // 进入新的作用域
        scc_vec_foreach(stmt->compound.block_items, i) {
            scc_ast_node_t *child_stmt =
                scc_vec_at(stmt->compound.block_items, i);
            if (SCC_AST_IS_A(scc_ast_stmt_t, child_stmt)) {
                ast_stmt_to_ir(ctx,
                               SCC_AST_CAST_TO(scc_ast_stmt_t, child_stmt));
            } else if (SCC_AST_IS_A(scc_ast_decl_t, child_stmt)) {
                ast_decl_to_ir(ctx,
                               SCC_AST_CAST_TO(scc_ast_decl_t, child_stmt));
            } else {
                UNREACHABLE();
            }
        }
        break;
    }

    case SCC_AST_STMT_EXPR: {
        ast_expr_to_ir(ctx, stmt->expr.expr);
        break;
    }

        // case SCC_AST_STMT_IF: {
        //     // 创建基本块 - 统一使用栈上分配
        //     scc_ir_bblock_t true_block;
        //     scc_ir_bblock_init(&true_block, "if_true");

        //     scc_ir_bblock_t false_block;
        //     scc_ir_bblock_init(&false_block, "if_false");

        //     scc_ir_bblock_t merge_block;
        //     scc_ir_bblock_init(&merge_block, "if_merge");

        //     // 转换条件
        //     scc_ir_node_ref_t cond_node = ast_expr_to_ir(ctx,
        //     stmt->if_stmt.cond);

        //     // 创建分支指令
        //     scc_ir_builder_branch();
        //     scc_ir_node_t br_node;
        //     init_branch_node(&br_node, cond, &true_block, &false_block);
        //     emit_instruction(ctx, &br_node);

        //     // 保存当前块
        //     ast2ir_ctx_t saved = ctx->current_ctx;

        //     // 生成true分支
        //     emit_basicblock(ctx, &true_block);
        //     ast_stmt_to_ir(ctx, stmt->if_stmt.then_stmt);

        //     // 跳转到合并块
        //     scc_ir_node_t jump_node;
        //     init_jump_node(&jump_node, &merge_block);
        //     emit_instruction(ctx, &jump_node);

        //     // 生成false分支（如果有）
        //     emit_basicblock(ctx, &false_block);
        //     if (stmt->if_stmt.opt_else_stmt) {
        //         ast_stmt_to_ir(ctx, stmt->if_stmt.opt_else_stmt);
        //     }

        //     // 跳转到合并块
        //     init_jump_node(&jump_node, &merge_block);
        //     emit_instruction(ctx, &jump_node);

        //     // 恢复上下文并设置当前块为合并块
        //     ctx->current_ctx = saved;
        //     emit_basicblock(ctx, &merge_block);
        //     break;
        // }

        // case SCC_AST_STMT_WHILE: {
        //     // 创建基本块 - 统一使用栈上分配
        //     scc_ir_bblock_t cond_block;
        //     scc_ir_bblock_init(&cond_block, "while_cond");

        //     scc_ir_bblock_t body_block;
        //     scc_ir_bblock_init(&body_block, "while_body");

        //     scc_ir_bblock_t exit_block;
        //     scc_ir_bblock_init(&exit_block, "while_exit");

        //     // 跳转到条件块
        //     scc_ir_node_t jump_node;
        //     init_jump_node(&jump_node, &cond_block);
        //     emit_instruction(ctx, &jump_node);

        //     // 保存当前块
        //     ast2ir_ctx_t saved = ctx->current_ctx;

        //     // 生成条件块
        //     emit_basicblock(ctx, &cond_block);
        //     scc_ir_node_t cond_node;
        //     ast_expr_to_ir(ctx, stmt->while_stmt.cond, &cond_node);

        //     // 将条件节点添加到程序节点列表
        //     scc_vec_push(ctx->program->global_vals, cond_node);
        //     scc_ir_node_t *cond =
        //         &scc_vec_at(ctx->program->global_vals,
        //                     scc_vec_size(ctx->program->global_vals) - 1);

        //     // 创建分支
        //     scc_ir_node_t br_node;
        //     init_branch_node(&br_node, cond, &body_block, &exit_block);
        //     emit_instruction(ctx, &br_node);

        //     // 生成循环体
        //     emit_basicblock(ctx, &body_block);
        //     ast_stmt_to_ir(ctx, stmt->while_stmt.body);

        //     // 跳转回条件块
        //     init_jump_node(&jump_node, &cond_block);
        //     emit_instruction(ctx, &jump_node);

        //     // 恢复上下文并设置当前块为退出块
        //     ctx->current_ctx = saved;
        //     emit_basicblock(ctx, &exit_block);
        //     break;
        // }

        // case SCC_AST_STMT_DO_WHILE: {
        //     // 创建基本块 - 统一使用栈上分配
        //     scc_ir_bblock_t cond_block;
        //     scc_ir_bblock_init(&cond_block, "do_while_cond");

        //     scc_ir_bblock_t body_block;
        //     scc_ir_bblock_init(&body_block, "do_while_body");

        //     scc_ir_bblock_t exit_block;
        //     scc_ir_bblock_init(&exit_block, "do_while_exit");

        //     // 跳转到循环体块
        //     scc_ir_node_t jump_node;
        //     init_jump_node(&jump_node, &body_block);
        //     emit_instruction(ctx, &jump_node);

        //     // 保存当前块
        //     ast2ir_ctx_t saved = ctx->current_ctx;

        //     // 生成循环体
        //     emit_basicblock(ctx, &body_block);
        //     ast_stmt_to_ir(ctx, stmt->do_while_stmt.body);

        //     // 跳转到条件块
        //     init_jump_node(&jump_node, &cond_block);
        //     emit_instruction(ctx, &jump_node);

        //     // 生成条件块
        //     emit_basicblock(ctx, &cond_block);
        //     scc_ir_node_t cond_node;
        //     ast_expr_to_ir(ctx, stmt->do_while_stmt.cond, &cond_node);

        //     // 将条件节点添加到程序节点列表
        //     scc_vec_push(ctx->program->global_vals, cond_node);
        //     scc_ir_node_t *cond =
        //         &scc_vec_at(ctx->program->global_vals,
        //                     scc_vec_size(ctx->program->global_vals) - 1);

        //     // 创建分支
        //     scc_ir_node_t br_node;
        //     init_branch_node(&br_node, cond, &body_block, &exit_block);
        //     emit_instruction(ctx, &br_node);

        //     // 恢复上下文并设置当前块为退出块
        //     ctx->current_ctx = saved;
        //     emit_basicblock(ctx, &exit_block);
        //     break;
        // }

        // case SCC_AST_STMT_FOR: {
        //     // 初始化部分
        //     if (SCC_AST_IS_A(scc_ast_expr_t, stmt->for_stmt.init)) {
        //         scc_ir_node_t dummy_node;
        //         ast_expr_to_ir(ctx,
        //                        SCC_AST_CAST_TO(scc_ast_expr_t,
        //                        stmt->for_stmt.init), &dummy_node);
        //     } else if (SCC_AST_IS_A(scc_ast_decl_t, stmt->for_stmt.init)) {
        //         ast_decl_to_ir(
        //             ctx, SCC_AST_CAST_TO(scc_ast_decl_t,
        //             stmt->for_stmt.init));
        //     } else {
        //         UNREACHABLE();
        //     }

        //     // 创建基本块 - 统一使用栈上分配
        //     scc_ir_bblock_t cond_block;
        //     scc_ir_bblock_init(&cond_block, "for_cond");

        //     scc_ir_bblock_t body_block;
        //     scc_ir_bblock_init(&body_block, "for_body");

        //     scc_ir_bblock_t exit_block;
        //     scc_ir_bblock_init(&exit_block, "for_exit");

        //     // 跳转到条件块
        //     scc_ir_node_t jump_node;
        //     init_jump_node(&jump_node, &cond_block);
        //     emit_instruction(ctx, &jump_node);

        //     // 保存当前块
        //     ast2ir_ctx_t saved = ctx->current_ctx;

        //     // 生成条件块
        //     emit_basicblock(ctx, &cond_block);

        //     scc_ir_node_t *cond = null;
        //     if (stmt->for_stmt.cond) {
        //         scc_ir_node_t cond_node;
        //         ast_expr_to_ir(ctx, stmt->for_stmt.cond, &cond_node);

        //         // 将条件节点添加到程序节点列表
        //         scc_vec_push(ctx->program->global_vals, cond_node);
        //         cond = &scc_vec_at(ctx->program->global_vals,
        //                            scc_vec_size(ctx->program->global_vals) -
        //                            1);
        //     } else {
        //         // 无条件循环，条件为真
        //         scc_ir_node_t true_node;
        //         init_const_int_node(&true_node, 1, null);

        //         // 将true节点添加到程序节点列表
        //         scc_vec_push(ctx->program->global_vals, true_node);
        //         cond = &scc_vec_at(ctx->program->global_vals,
        //                            scc_vec_size(ctx->program->global_vals) -
        //                            1);
        //     }

        //     // 创建分支
        //     scc_ir_node_t br_node;
        //     init_branch_node(&br_node, cond, &body_block, &exit_block);
        //     emit_instruction(ctx, &br_node);

        //     // 生成循环体
        //     emit_basicblock(ctx, &body_block);
        //     ast_stmt_to_ir(ctx, stmt->for_stmt.body);

        //     // 执行迭代表达式（如果存在）
        //     if (stmt->for_stmt.iter) {
        //         scc_ir_node_t dummy_node;
        //         ast_expr_to_ir(ctx, stmt->for_stmt.iter, &dummy_node);
        //     }

        //     // 跳转回条件块
        //     init_jump_node(&jump_node, &cond_block);
        //     emit_instruction(ctx, &jump_node);

        //     // 恢复上下文并设置当前块为退出块
        //     ctx->current_ctx = saved;
        //     emit_basicblock(ctx, &exit_block);
        //     break;
        // }

    case SCC_AST_STMT_RETURN: {
        if (stmt->return_stmt.expr) {
            scc_ir_node_ref_t ret_val_node =
                ast_expr_to_ir(ctx, stmt->return_stmt.expr);

            scc_ir_builder_ret(ctx, ret_val_node);
        } else {
            scc_ir_builder_ret_void(ctx);
        }
        break;
    }

    default:
        LOG_WARN("Unsupported statement type: %d", stmt->base.type);
        break;
    }
}

// 转换AST声明为IR
static void ast_decl_to_ir(scc_ir_builder_t *ctx, scc_ast_decl_t *decl) {
    if (ctx == null || decl == null) {
        LOG_ERROR("Invalid argument");
        return;
    }

    switch (decl->base.type) {
    case SCC_AST_DECL_VAR: {
        // 转换类型
        scc_ir_type_ref_t ir_type = ast_type_to_ir_type(ctx, decl->var.type);

        // 创建分配节点
        scc_ir_node_ref_t alloc_val_node =
            scc_ir_builder_alloca(ctx, ir_type, decl->var.name);

        // 如果有初始化表达式
        if (!decl->var.init) {
            break;
        }
        scc_ir_node_ref_t init_val_node = ast_expr_to_ir(ctx, decl->var.init);

        // 将初始化值节点添加到程序节点列表
        // scc_vec_push(ctx->program->global_vals, init_val_node);
        // scc_ir_node_t *init_val =
        //     &scc_vec_at(ctx->program->global_vals,
        //                 scc_vec_size(ctx->program->global_vals) - 1);

        scc_ir_builder_store(ctx, alloc_val_node, init_val_node);
        break;
    }

    case SCC_AST_DECL_FUNC: {
        // TODO params name
        scc_ir_type_ref_t func_type = ast_type_to_ir_type(ctx, decl->func.type);
        scc_ir_builder_begin_func(ctx, decl->func.name, func_type, null);
        // 处理函数体（如果有）
        if (decl->func.body) {
            scc_ir_builder_begin_bblock(ctx, "entry");
            ast_stmt_to_ir(ctx, decl->func.body);
            scc_ir_builder_end_bblock(ctx);
        }
        scc_ir_builder_end_func(ctx);
        break;
    }

    default:
        LOG_WARN("Unsupported declaration type: %d", decl->base.type);
        break;
    }
}

// 主转换函数：将AST翻译单元转换为IR程序
void scc_ast2ir(scc_ast_translation_unit_t *tu, scc_ir_builder_t *builder) {
    if (tu == null || builder == null) {
        LOG_ERROR("Invalid argument");
        return;
    }

    // 初始化上下文
    scc_ir_builder_init(builder);

    // 转换所有声明
    scc_vec_foreach(tu->declarations, i) {
        scc_ast_decl_t *decl = scc_vec_at(tu->declarations, i);
        ast_decl_to_ir(builder, decl);
    }
}