#include "irgen/IRGen.h"

#include <stdexcept>
#include <string>

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

// 辅助：求值表达式为整数（用于数组维度）
static int EvalExprInt(SysYParser::ExpContext* ctx) {
  if (!ctx) return 0;
  try {
    auto cv = sem::EvaluateExp(*ctx->addExp());
    return static_cast<int>(cv.int_val);
  } catch (...) {
    return 0;
  }
}

// ─── 构造函数 ─────────────────────────────────────────────────────────────────
IRGenImpl::IRGenImpl(ir::Module& module, const SemanticContext& sema)
    : module_(module),
      sema_(sema),
      func_(nullptr),
      builder_(module.GetContext(), nullptr) {}

// ─── visitCompUnit ────────────────────────────────────────────────────────────
std::any IRGenImpl::visitCompUnit(SysYParser::CompUnitContext* ctx) {
  if (!ctx) {
    throw std::runtime_error(FormatError("irgen", "缺少编译单元"));
  }
  // 先处理全局声明
  in_global_scope_ = true;
  for (auto* decl : ctx->decl()) {
    if (decl) decl->accept(this);
  }
  in_global_scope_ = false;
  // 再生成函数
  for (auto* funcDef : ctx->funcDef()) {
    if (funcDef) funcDef->accept(this);
  }
  return {};
}

// ─── visitFuncDef ─────────────────────────────────────────────────────────────
std::any IRGenImpl::visitFuncDef(SysYParser::FuncDefContext* ctx) {
  if (!ctx || !ctx->Ident()) {
    throw std::runtime_error(FormatError("irgen", "缺少函数定义或函数名"));
  }
  if (!ctx->block()) {
    throw std::runtime_error(FormatError("irgen", "函数体为空"));
  }

  std::string func_name = ctx->Ident()->getText();

  // 确定返回类型
  std::shared_ptr<ir::Type> ret_type;
  if (!ctx->funcType()) {
    throw std::runtime_error(FormatError("irgen", "缺少函数返回类型"));
  }
  if (ctx->funcType()->Void()) {
    ret_type = ir::Type::GetVoidType();
  } else if (ctx->funcType()->Int()) {
    ret_type = ir::Type::GetInt32Type();
  } else if (ctx->funcType()->Float()) {
    ret_type = ir::Type::GetFloat32Type();
  } else {
    throw std::runtime_error(
        FormatError("irgen", "函数 " + func_name + " 返回类型不支持"));
  }

  func_ = module_.CreateFunction(func_name, ret_type);
  storage_map_.clear();

  // 设置插入点到入口块
  builder_.SetInsertPoint(func_->GetEntry());
  builder_.SetAllocaBlock(func_->GetEntry());

  // 处理参数
  if (ctx->funcFParams()) {
    for (auto* param : ctx->funcFParams()->funcFParam()) {
      if (!param || !param->Ident()) continue;
      std::string pname = param->Ident()->getText();

      bool is_float = param->bType() && param->bType()->Float();
      bool is_array_param = !param->L_BRAKT().empty();

      if (is_array_param) {
        // 数组参数：传递为指针
        auto ptr_type = is_float ? ir::Type::GetPtrFloat32Type() : ir::Type::GetPtrInt32Type();
        ir::Argument* arg = func_->AddArgument(ptr_type, pname);
        storage_map_[param] = arg;

        // 记录维度信息（第一维未知，后续维度从 exp 中获取）
        std::vector<int> dims;
        dims.push_back(-1);  // 第一维未知
        for (auto* exp_ctx : param->exp()) {
          dims.push_back(EvalExprInt(exp_ctx));
        }
        array_dims_[param] = dims;  // 始终记录，包括一维数组参数
      } else {
        // 标量参数
        auto arg_type = is_float ? ir::Type::GetFloat32Type() : ir::Type::GetInt32Type();
        ir::Argument* arg = func_->AddArgument(arg_type, pname);
        auto* slot = is_float ? builder_.CreateAllocaF32(pname + ".addr")
                              : builder_.CreateAllocaI32(pname + ".addr");
        storage_map_[param] = slot;
        builder_.CreateStore(arg, slot);
      }
    }
  }

  // 生成函数体
  ctx->block()->accept(this);

  // 若最后一个基本块没有 terminator，自动补 ret
  auto* last_bb = builder_.GetInsertBlock();
  if (last_bb && !last_bb->HasTerminator()) {
    if (ret_type->IsVoid()) {
      builder_.CreateRetVoid();
    } else if (ret_type->IsFloat32()) {
      builder_.CreateRet(builder_.CreateConstFloat(0.0f));
    } else {
      builder_.CreateRet(builder_.CreateConstInt(0));
    }
  }

  return {};
}