#include "irgen/IRGen.h"

#include <stdexcept>

#include "SysYParser.h"
#include "ir/IR.h"
#include "utils/Log.h"
#include <functional>

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, IRGenContext& sema)
    : module_(module),
      sema_(sema),
      func_(nullptr),
      builder_(module.GetContext(), nullptr) {}

// 编译单元的 IR 生成当前只实现了最小功能：
// - Module 已在 GenerateIR 中创建，这里只负责继续生成其中的内容；
// - 当前会读取编译单元中的函数定义，并交给 visitFuncDef 生成函数 IR；
//
// 当前还没有实现：
// - 多个函数定义的遍历与生成；
// - 全局变量、全局常量的 IR 生成。
std::any IRGenImpl::visitCompUnit(SysYParser::CompUnitContext* ctx) {
  if (!ctx) {
    throw std::runtime_error(FormatError("irgen", "缺少编译单元"));
  }
  // for simplicity take first function definition
  if (ctx->funcDef().empty()) {
    throw std::runtime_error(FormatError("irgen", "缺少函数定义"));
  }
  // 先处理顶层声明（仅支持简单的 const int 初始化作为全局常量）
  for (auto* decl : ctx->decl()) {
    if (!decl) continue;
    if (decl->constDecl()) {
      auto* cdecl = decl->constDecl();
      for (auto* cdef : cdecl->constDef()) {
        if (!cdef || !cdef->Ident() || !cdef->constInitVal() || !cdef->constInitVal()->constExp()) continue;
        // 仅支持形如: const int a = 10; 的简单常量初始化（字面量）
        auto* add = cdef->constInitVal()->constExp()->addExp();
        if (!add) continue;
        try {
          int v = std::stoi(add->getText());
          auto* cval = module_.GetContext().GetConstInt(v);
          name_map_[cdef->Ident()->getText()] = cval;
        } catch (...) {
          // 无法解析则跳过，全局复杂常量暂不支持
        }
      }
    }
    // 支持简单的全局变量声明（数组或标量），初始化为零
    if (decl->varDecl()) {
      auto* vdecl = decl->varDecl();
      if (!vdecl->bType()) continue;
      std::string btype = vdecl->bType()->getText();
      for (auto* vdef : vdecl->varDef()) {
        if (!vdef) continue;
        LogInfo(std::string("[irgen] global varDef text=") + vdef->getText() + std::string(" ident=") + (vdef->Ident() ? vdef->Ident()->getText() : std::string("<none>")) + std::string(" dims=") + std::to_string((int)vdef->constExp().size()), std::cerr);
        if (!vdef || !vdef->Ident()) continue;
        std::string name = vdef->Ident()->getText();
        // array globals
        if (!vdef->constExp().empty()) {
          std::vector<int> dims;
          bool ok = true;
          for (auto* ce : vdef->constExp()) {
            try {
              int val = 0;
              auto anyv = sema_.GetConstVal(ce);
              if (anyv.has_value()) {
                if (anyv.type() == typeid(int)) val = std::any_cast<int>(anyv);
                else if (anyv.type() == typeid(long)) val = (int)std::any_cast<long>(anyv);
                else throw std::runtime_error("not-const-int");
              } else {
                // try literal parse
                try {
                  val = std::stoi(ce->addExp()->getText());
                } catch (...) {
                  // try lookup in name_map_ for previously created const
                  std::string t = ce->addExp()->getText();
                  auto it = name_map_.find(t);
                  if (it != name_map_.end() && it->second && it->second->IsConstant()) {
                    if (auto* ci = dynamic_cast<ir::ConstantInt*>(it->second)) {
                      val = ci->GetValue();
                    } else {
                      ok = false; break;
                    }
                  } else {
                    ok = false; break;
                  }
                }
              }
              dims.push_back(val);
            } catch (...) { ok = false; break; }
          }
          if (!ok) continue;
          // build zero constant array similar to visitVarDef
          std::function<ir::ConstantValue*(const std::vector<int>&, size_t, std::shared_ptr<ir::Type>)> buildZero;
          buildZero = [&](const std::vector<int>& ds, size_t idx, std::shared_ptr<ir::Type> elemTy) -> ir::ConstantValue* {
            if (idx >= ds.size()) return nullptr;
            std::vector<ir::ConstantValue*> elems;
            if (idx + 1 == ds.size()) {
              for (int i = 0; i < ds[idx]; ++i) {
                if (elemTy->IsFloat32()) elems.push_back(module_.GetContext().GetConstFloat(0.0f));
                else elems.push_back(module_.GetContext().GetConstInt(0));
              }
            } else {
              for (int i = 0; i < ds[idx]; ++i) {
                ir::ConstantValue* sub = buildZero(ds, idx + 1, elemTy);
                if (sub) elems.push_back(sub);
                else elems.push_back(module_.GetContext().GetConstInt(0));
              }
            }
            std::function<std::shared_ptr<ir::Type>(size_t)> makeArrayType = [&](size_t level) -> std::shared_ptr<ir::Type> {
              if (level + 1 >= ds.size()) return ir::Type::GetArrayType(elemTy, ds[level]);
              auto sub = makeArrayType(level + 1);
              return ir::Type::GetArrayType(sub, ds[level]);
            };
            auto at_real = makeArrayType(idx);
            return new ir::ConstantArray(at_real, elems);
          };
          std::shared_ptr<ir::Type> elemTy = (btype == "float") ? ir::Type::GetFloat32Type() : ir::Type::GetInt32Type();
          ir::ConstantValue* zero = buildZero(dims, 0, elemTy);
          auto gvty = ir::Type::GetPointerType(zero ? zero->GetType() : ir::Type::GetPointerType(elemTy));
          ir::GlobalValue* gv = module_.CreateGlobalVariable(name, gvty, zero);
          name_map_[name] = gv;
          LogInfo(std::string("[irgen] created global ") + name, std::cerr);
        } else {
          // scalar global
          std::shared_ptr<ir::Type> elemTy = (btype == "float") ? ir::Type::GetFloat32Type() : ir::Type::GetInt32Type();
          ir::ConstantValue* init = nullptr;
          if (btype == "float") init = module_.GetContext().GetConstFloat(0.0f);
          else init = module_.GetContext().GetConstInt(0);
          ir::GlobalValue* gv = module_.CreateGlobalVariable(name, ir::Type::GetPointerType(elemTy), init);
          name_map_[name] = gv;
          LogInfo(std::string("[irgen] created global ") + name, std::cerr);
        }
      }
    }
  }

  // 生成编译单元中所有函数定义（之前只生成第一个函数）
  for (size_t i = 0; i < ctx->funcDef().size(); ++i) {
    if (ctx->funcDef(i)) ctx->funcDef(i)->accept(this);
  }
  return {};
}

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

// 因此这里目前只支持最小的“无参 int 函数”生成。
// 因此这里目前只支持最小的“无参 int 函数”生成。
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->Ident()) {
    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()->getText() == "int") ret_type = ir::Type::GetInt32Type();
  else if (ctx->funcType()->getText() == "float") ret_type = ir::Type::GetFloat32Type();
  else if (ctx->funcType()->getText() == "void") ret_type = ir::Type::GetVoidType();
  else throw std::runtime_error(FormatError("irgen", "当前仅支持 int/float/void 函数"));

  std::vector<std::shared_ptr<ir::Type>> param_types;
  if (ctx->funcFParams()) {
    for (auto* param : ctx->funcFParams()->funcFParam()) {
      if (param->bType()->getText() == "int") param_types.push_back(ir::Type::GetInt32Type());
      else if (param->bType()->getText() == "float") param_types.push_back(ir::Type::GetFloat32Type());
      else throw std::runtime_error(FormatError("irgen", "当前仅支持 int/float 参数"));
    }
  }

  func_ = module_.CreateFunction(ctx->Ident()->getText(), ret_type, param_types);
  builder_.SetInsertPoint(func_->GetEntry());

  // Allocate storage for parameters
  if (ctx->funcFParams()) {
    int idx = 0;
    for (auto* param : ctx->funcFParams()->funcFParam()) {
      std::string param_name = param->Ident()->getText();
      ir::AllocaInst* alloca = nullptr;
      if (param->bType()->getText() == "float") alloca = builder_.CreateAllocaFloat(param_name);
      else alloca = builder_.CreateAllocaI32(param_name);
      storage_map_[param] = alloca;
      name_map_[param_name] = alloca;
      // Store the argument value
      auto* arg = func_->GetParams()[idx];
      builder_.CreateStore(arg, alloca);
      idx++;
    }
  }


  ctx->block()->accept(this);
  // 如果函数体末尾没有显式终结（如 void 函数没有 return），补一个隐式 return
  if (!builder_.GetInsertBlock()->HasTerminator()) {
    builder_.CreateRet(nullptr);
  }
  // 语义正确性主要由 sema 保证，这里只兜底检查 IR 结构是否合法。
  VerifyFunctionStructure(*func_);
  return {};
}
