#include "irgen/IRGen.h"

#include <stdexcept>

#include "antlr4/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) {
  // 初始化作用域栈，至少有一个全局作用域（但当前未使用全局变量）
  scope_storage_.emplace_back();
}

ir::AllocaInst* IRGenImpl::CreateEntryAllocaI32(const std::string& name) {
  if (!func_) {
    throw std::runtime_error(FormatError("irgen", "局部 alloca 必须位于函数内"));
  }
  auto* saved = builder_.GetInsertBlock();
  builder_.SetInsertPoint(func_->GetEntry());
  auto* slot = builder_.CreateAllocaI32(name);
  builder_.SetInsertPoint(saved);
  return slot;
}

ir::AllocaInst* IRGenImpl::CreateEntryAllocaArray(int count, const std::string& name) {
  if (!func_) {
    throw std::runtime_error(FormatError("irgen", "局部 alloca 必须位于函数内"));
  }
  auto* saved = builder_.GetInsertBlock();
  builder_.SetInsertPoint(func_->GetEntry());
  auto* slot = builder_.CreateAllocaArray(count, name);
  builder_.SetInsertPoint(saved);
  return slot;
}

ir::AllocaInst* IRGenImpl::CreateEntryAllocaF32(const std::string& name) {
  if (!func_) {
    throw std::runtime_error(FormatError("irgen", "局部 alloca 必须位于函数内"));
  }
  auto* saved = builder_.GetInsertBlock();
  builder_.SetInsertPoint(func_->GetEntry());
  auto* slot = builder_.CreateAllocaF32(name);
  builder_.SetInsertPoint(saved);
  return slot;
}

ir::AllocaInst* IRGenImpl::CreateEntryAllocaF32Array(int count, const std::string& name) {
  if (!func_) {
    throw std::runtime_error(FormatError("irgen", "局部 alloca 必须位于函数内"));
  }
  auto* saved = builder_.GetInsertBlock();
  builder_.SetInsertPoint(func_->GetEntry());
  auto* slot = builder_.CreateAllocaF32Array(count, name);
  builder_.SetInsertPoint(saved);
  return slot;
}

// 预声明 SysY 运行时外部函数（putint / putch / getint / getch 等）。
void IRGenImpl::DeclareRuntimeFunctions() {
  auto i32 = ir::Type::GetInt32Type();
  auto void_ = ir::Type::GetVoidType();

  auto decl = [&](const std::string& name,
                   std::shared_ptr<ir::Type> ret,
                   std::vector<std::shared_ptr<ir::Type>> params) {
    if (!module_.FindFunction(name)) {
      auto* f = module_.CreateFunction(name, ret, params);
      f->SetExternal(true);
    }
  };

  // 整数 I/O
  decl("getint",  i32,   {});
  decl("getch",   i32,   {});
  decl("putint",  void_, {i32});
  decl("putch",   void_, {i32});
  // 数组 I/O
  decl("getarray",  i32,   {ir::Type::GetPtrInt32Type()});
  decl("putarray",  void_, {i32, ir::Type::GetPtrInt32Type()});
  // 浮点 I/O
  decl("getfloat",  ir::Type::GetFloat32Type(), {});
  decl("getfarray", i32, {ir::Type::GetPtrFloat32Type()});
  decl("putfloat",  void_, {ir::Type::GetFloat32Type()});
  decl("putfarray", void_, {i32, ir::Type::GetPtrFloat32Type()});
  // 时间
  decl("starttime", void_, {});
  decl("stoptime",  void_, {});
  // 通用内存清零（用于局部 float 大数组初始化）
  decl("memset", void_, {ir::Type::GetPtrFloat32Type(), i32, i32});
}

// 编译单元 IR 生成：
// 1. 预声明 SysY runtime；
// 2. 处理全局变量/常量声明；
// 3. 生成各函数 IR。
std::any IRGenImpl::visitCompUnit(SysYParser::CompUnitContext* ctx) {
  if (!ctx) {
    throw std::runtime_error(FormatError("irgen", "缺少编译单元"));
  }

  DeclareRuntimeFunctions();

  // 全局声明（func_ == nullptr 时 visitVarDef/visitConstDef 会走全局路径）
  for (auto* decl : ctx->decl()) {
    if (decl) decl->accept(this);
  }

  for (auto* func : ctx->funcDef()) {
    if (func) func->accept(this);
  }
  return {};
}

// 函数 IR 生成当前实现了：
// 1. 获取函数名；
// 2. 支持 int 与 void 返回类型；
// 3. 支持 int 形参：入口处为每个参数 alloca + store；
// 4. 在 Module 中创建 Function；
// 5. 将 builder 插入点设置到入口基本块；
// 6. 继续生成函数体。
//
// 当前还没有实现：
// - float 参数/返回类型；
// - 数组类型形参；
// - FunctionType 这样的函数类型对象（参数类型目前只用 shared_ptr<Type>）。

std::any IRGenImpl::visitFuncDef(SysYParser::FuncDefContext* ctx) {
  if (!ctx) {
    throw std::runtime_error(FormatError("irgen", "缺少函数定义"));
  }
  if (!ctx->blockStmt()) {
    throw std::runtime_error(FormatError("irgen", "函数体为空"));
  }
  if (!ctx->ID()) {
    throw std::runtime_error(FormatError("irgen", "缺少函数名"));
  }
  if (!ctx->funcType()) {
    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()->VOID()) {
    ret_type = ir::Type::GetVoidType();
  } else if (ctx->funcType()->FLOAT()) {
    ret_type = ir::Type::GetFloat32Type();
  } else {
    throw std::runtime_error(FormatError("irgen", "当前仅支持 int/void/float 返回类型"));
  }

  // 收集形参类型（支持 int 标量和 int 数组参数）。
  std::vector<std::shared_ptr<ir::Type>> param_types;
  std::vector<std::string> param_names;
  std::vector<bool> param_is_array;

  if (auto* fparams = ctx->funcFParams()) {
    for (auto* fp : fparams->funcFParam()) {
      if (!fp || !fp->btype()) {
        throw std::runtime_error(
            FormatError("irgen", "缺少参数类型"));
      }
      bool is_int = fp->btype()->INT() != nullptr;
      bool is_float = fp->btype()->FLOAT() != nullptr;
      if (!is_int && !is_float) {
        throw std::runtime_error(
            FormatError("irgen", "当前仅支持 int/float 类型形参"));
      }
      bool is_arr = !fp->LBRACK().empty();
      param_is_array.push_back(is_arr);
      if (is_arr) {
        param_types.push_back(is_int ? ir::Type::GetPtrInt32Type()
                                     : ir::Type::GetPtrFloat32Type());
      } else {
        param_types.push_back(is_int ? ir::Type::GetInt32Type()
                                     : ir::Type::GetFloat32Type());
      }
      param_names.push_back(fp->ID() ? fp->ID()->getText() : "");
    }
  }

  func_ = module_.CreateFunction(ctx->ID()->getText(), ret_type, param_types);
  auto* body_entry = func_->CreateBlock(NextBlockName());
  builder_.SetInsertPoint(body_entry);
  storage_map_.clear();
  named_storage_.clear();
  local_array_dims_.clear();

  // 第二遍：处理形参（现在有插入点，可以生成 alloca 等）
  auto* fparams = ctx->funcFParams();
  for (size_t i = 0; i < param_names.size(); ++i) {
    auto* arg = func_->GetArgument(i);
    if (param_is_array[i]) {
      // 数组参数：直接存入 named_storage_，维度用 EvalExpAsConst 获取
      if (!param_names[i].empty()) {
        named_storage_[param_names[i]] = arg;
        std::vector<int> dims = {-1};  // 首维未知
        if (fparams) {
          auto fp_list = fparams->funcFParam();
          if (i < fp_list.size()) {
            for (auto* dim_exp : fp_list[i]->exp()) {
              dims.push_back(EvalExpAsConst(dim_exp));
            }
          }
        }
        local_array_dims_[param_names[i]] = dims;
      }
    } else {
      // 标量参数：alloca + store
      ir::AllocaInst* slot = nullptr;
      if (arg->GetType()->IsFloat32()) {
        slot = CreateEntryAllocaF32(module_.GetContext().NextTemp());
      } else {
        slot = CreateEntryAllocaI32(module_.GetContext().NextTemp());
      }
      builder_.CreateStore(arg, slot);
      if (!param_names[i].empty()) {
        named_storage_[param_names[i]] = slot;
      }
    }
  }

  short_circuit_slot_ = CreateEntryAllocaI32(module_.GetContext().NextTemp());

  // 进入函数参数作用域
  scope_storage_.emplace_back();

  // 处理参数：为每个参数创建 alloca 并放入符号表
  for (const auto& param_name : param_names) {
    auto* alloca = builder_.CreateAllocaI32(module_.GetContext().NextTemp());
    param_slots_.push_back(alloca);
    scope_storage_.back()[param_name] = alloca;
    named_storage_[param_name] = alloca;
  }

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

  // 入口块只用于静态栈槽分配，末尾统一跳到函数体起始块。
  auto* saved = builder_.GetInsertBlock();
  builder_.SetInsertPoint(func_->GetEntry());
  if (!func_->GetEntry()->HasTerminator()) {
    builder_.CreateBr(body_entry);
  }
  builder_.SetInsertPoint(saved);

  // 对于 void 函数，若末尾块无 terminator，自动补 ret void。
  if (ret_type->IsVoid()) {
    auto* bb = builder_.GetInsertBlock();
    if (bb && !bb->HasTerminator()) {
      builder_.CreateRetVoid();
    }
  }

  // 语义正确性主要由 sema 保证，这里只兜底检查 IR 结构是否合法。
  VerifyFunctionStructure(*func_);
  short_circuit_slot_ = nullptr;
  func_ = nullptr;  // 回到全局作用域
  return {};
}