#include "irgen/IRGen.h"

#include <stdexcept>

#include "SysYParser.h"
#include "ir/IR.h"
#include "utils/Log.h"

namespace {

void VerifyFunctionStructure(const ir::Function& func) {
  // 当前 IRGen 仍是单入口、顺序生成；这里在生成结束后补一层块终结校验。
  for (const auto& bb : func.GetBlocks()) {
    if (!bb || !bb->HasTerminator()) {
      throw std::runtime_error(
          FormatError("irgen", "基本块未正确终结: " +
                                   (bb ? bb->GetName() : std::string("<null>"))));
    }
  }
}

}  // namespace

IRGenImpl::IRGenImpl(ir::Module& module, const SemanticContext& sema)
    : module_(module),
      sema_(sema),
      func_(nullptr),
      builder_(module.GetContext(), nullptr) {}

// 编译单元的 IR 生成当前只实现了最小功能：
// - Module 已在 GenerateIR 中创建，这里只负责继续生成其中的内容；
// - 当前会读取编译单元中的 topLevelItem，找到 funcDef 后生成函数 IR；
//
// 当前还没有实现：
// - 多个函数定义的遍历与生成；
// - 全局变量、全局常量的 IR 生成。
std::any IRGenImpl::visitCompUnit(SysYParser::CompUnitContext* ctx) {
  if (!ctx) {
    throw std::runtime_error(FormatError("irgen", "缺少编译单元"));
  }
  // 初始化全局作用域
  storage_map_stack_.push_back({});
  const_values_stack_.push_back({});

  // 遍历所有 topLevelItem
  for (auto* item : ctx->topLevelItem()) {
    if (!item) continue;
    if (item->funcDef()) {
      item->funcDef()->accept(this);
    } else if (item->decl()) {
      item->decl()->accept(this);
    }
  }

  // 退出全局作用域
  storage_map_stack_.pop_back();
  const_values_stack_.pop_back();

  return {};
}

// 函数 IR 生成当前实现了：
// 1. 获取函数名；
// 2. 检查函数返回类型；
// 3. 在 Module 中创建 Function；
// 4. 将 builder 插入点设置到入口基本块；
// 5. 继续生成函数体。
//
// 当前还没有实现：
// - 通用函数返回类型处理；
// - 形参列表遍历与参数类型收集；
// - FunctionType 这样的函数类型对象；
// - Argument/形式参数 IR 对象；
// - 入口块中的参数初始化逻辑。
// ...

std::any IRGenImpl::visitFuncDef(SysYParser::FuncDefContext* ctx) {
  if (!ctx) {
    throw std::runtime_error(FormatError("irgen", "缺少函数定义"));
  }
  if (!ctx->block()) {
    throw std::runtime_error(FormatError("irgen", "函数体为空"));
  }
  if (!ctx->ID()) {
    throw std::runtime_error(FormatError("irgen", "缺少函数名"));
  }
  
  std::shared_ptr<ir::Type> ret_type;
  if (ctx->funcType()->INT()) {
    ret_type = ir::Type::GetInt32Type();
  } else if (ctx->funcType()->FLOAT()) {
    ret_type = ir::Type::GetFloatType();
  } else if (ctx->funcType()->VOID()) {
    ret_type = ir::Type::GetVoidType();
  } else {
    throw std::runtime_error(FormatError("irgen", "未知的函数返回类型"));
  }

  func_ = module_.CreateFunction(ctx->ID()->getText(), ret_type);
  ir::BasicBlock* alloca_bb = func_->CreateBlock("alloca");
  ir::BasicBlock* entry_bb = func_->CreateBlock("entry");
  builder_.SetInsertPoint(entry_bb);
  
  // 进入函数作用域，压入一个新的 map
  storage_map_stack_.push_back({});
  const_values_stack_.push_back({});

  if (ctx->funcFParams()) {
    for (auto* paramCtx : ctx->funcFParams()->funcFParam()) {
      std::shared_ptr<ir::Type> param_type;
      bool is_array = !paramCtx->LBRACK().empty();
      
      auto base_sema_ty = paramCtx->bType()->INT() ? SemanticType::Int : SemanticType::Float;
      auto base_ir_ty = (base_sema_ty == SemanticType::Int) ? ir::Type::GetInt32Type() : ir::Type::GetFloatType();

      if (is_array) {
        std::shared_ptr<ir::Type> elem_ty = base_ir_ty;
        auto exps = paramCtx->exp();
        for (auto it = exps.rbegin(); it != exps.rend(); ++it) {
          int dim = EvaluateConstInt(*it);
          elem_ty = ir::Type::GetArrayType(elem_ty, dim);
        }
        param_type = ir::Type::GetPointerType(elem_ty);
      } else {
        param_type = base_ir_ty;
      }

      std::string arg_name = paramCtx->ID()->getText();
      auto* arg = func_->AddArgument(param_type, "%arg" + std::to_string(func_->GetArgs().size()));
      
      // Ensure param alloca is in alloca_bb
      auto* current_bb = builder_.GetInsertBlock();
      builder_.SetInsertPoint(alloca_bb);
      ir::Instruction* alloca_inst = builder_.CreateAlloca(param_type, module_.GetContext().NextTemp());
      builder_.SetInsertPoint(current_bb);
      
      builder_.CreateStore(arg, alloca_inst);
      storage_map_stack_.back()[arg_name] = alloca_inst;
    }
  }

  ctx->block()->accept(this);
  
  // Implicit return for void functions or main
  if (!builder_.GetInsertBlock()->HasTerminator()) {
    if (func_->GetType()->IsVoid()) {
      builder_.CreateRet(nullptr);
    } else if (func_->GetName() == "main") {
      builder_.CreateRet(builder_.CreateConstInt(0));
    } else {
      if (func_->GetType()->IsFloat()) {
        builder_.CreateRet(module_.GetContext().GetConstFloat(0.0f));
      } else {
        builder_.CreateRet(builder_.CreateConstInt(0));
      }
    }
  }
  
  // Branch from alloca_bb to entry_bb
  builder_.SetInsertPoint(alloca_bb);
  builder_.CreateBr(entry_bb);

  VerifyFunctionStructure(*func_);
  func_ = nullptr;
  
  // 退出函数作用域，弹出 map
  storage_map_stack_.pop_back();
  const_values_stack_.pop_back();
  
  return {};
}
