#include "irgen/IRGen.h"

#include <memory>
#include <stdexcept>
#include <vector>

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

namespace {

void VerifyFunctionStructure(const ir::Function& func) {
  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) {}

// 修正：没有 mainFuncDef，通过函数名找到 main
std::any IRGenImpl::visitCompUnit(SysYParser::CompUnitContext* ctx) {
  std::cout << "[DEBUG IRGEN] visitCompUnit" << std::endl;
  if (!ctx) {
    throw std::runtime_error(FormatError("irgen", "缺少编译单元"));
  }

  // 处理全局变量声明
  for (auto* decl : ctx->decl()) {
    if (decl) {
      decl->accept(this);
    }
  }

  // 处理所有函数定义
  for (auto* funcDef : ctx->funcDef()) {
    if (funcDef) {
      funcDef->accept(this);
    }
  }

  return {};
}

std::any IRGenImpl::visitFuncDef(SysYParser::FuncDefContext* ctx) {
  std::cout << "[DEBUG IRGEN] visitFuncDef: " << (ctx && ctx->Ident() ? ctx->Ident()->getText() : "<null>") << std::endl;
  if (!ctx) {
    throw std::runtime_error(FormatError("irgen", "缺少函数定义"));
  }

  if (!ctx->Ident()) {
    throw std::runtime_error(FormatError("irgen", "缺少函数名"));
  }

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

  if (!ctx->block()) {
    throw std::runtime_error(FormatError("irgen", "函数体为空"));
  }

  std::shared_ptr<ir::Type> ret_type = ir::Type::GetInt32Type();
  if (ctx->funcType()) {
    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::GetFloatType();
    }
  }

  std::vector<std::shared_ptr<ir::Type>> param_types;
  if (ctx->funcFParams()) {
    for (auto* param : ctx->funcFParams()->funcFParam()) {
      if (!param || !param->bType()) continue;
      std::shared_ptr<ir::Type> param_ty;
      if (param->bType()->Int()) {
        param_ty = ir::Type::GetInt32Type();
      } else if (param->bType()->Float()) {
        param_ty = ir::Type::GetFloatType();
      } else {
        param_ty = ir::Type::GetInt32Type();
      }
      if (!param->L_BRACK().empty()) {
        if (param_ty->IsInt32()) {
          param_ty = ir::Type::GetPtrInt32Type();
        } else if (param_ty->IsFloat()) {
          param_ty = ir::Type::GetPtrFloatType();
        }
      }
      param_types.push_back(param_ty);
    }
  }

  auto func_type = ir::Type::GetFunctionType(ret_type, param_types);
  func_ = module_.CreateFunction(funcName, func_type);
  builder_.SetInsertPoint(func_->GetEntry());
  storage_map_.clear();
  param_map_.clear();

  // 函数参数 (按照语义分析、symbol table 定义顺序)
  if (ctx->funcFParams()) {
    for (auto* param : ctx->funcFParams()->funcFParam()) {
      if (!param || !param->Ident()) continue;
      std::string name = param->Ident()->getText();
      std::shared_ptr<ir::Type> param_ty;
      if (param->bType()->Int()) {
        param_ty = ir::Type::GetInt32Type();
      } else if (param->bType()->Float()) {
        param_ty = ir::Type::GetFloatType();
      } else {
        param_ty = ir::Type::GetInt32Type();
      }
      if (!param->L_BRACK().empty()) {
        if (param_ty->IsInt32()) {
          param_ty = ir::Type::GetPtrInt32Type();
        } else if (param_ty->IsFloat()) {
          param_ty = ir::Type::GetPtrFloatType();
        }
      }
      auto* arg = func_->AddArgument(std::make_unique<ir::Argument>(param_ty, name));
      ir::AllocaInst* slot;
      if (param_ty->IsInt32() || param_ty->IsPtrInt32()) {
        slot = builder_.CreateAllocaI32(module_.GetContext().NextTemp());
      } else if (param_ty->IsFloat() || param_ty->IsPtrFloat()) {
        slot = builder_.CreateAllocaFloat(module_.GetContext().NextTemp());
      } else {
        throw std::runtime_error(FormatError("irgen", "不支持的参数类型"));
      }
      builder_.CreateStore(arg, slot);
      param_map_[name] = slot;
    }
  }

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

  if (!func_->GetEntry()->HasTerminator()) {
    auto retVal = builder_.CreateConstInt(0);
    builder_.CreateRet(retVal);
  }

  VerifyFunctionStructure(*func_);
  func_ = nullptr;
  return {};
}

std::any IRGenImpl::visitBlock(SysYParser::BlockContext* ctx) {
  std::cout << "[DEBUG IRGEN] visitBlock: " << (ctx ? ctx->getText() : "<null>") << std::endl;
  if (!ctx) {
    throw std::runtime_error(FormatError("irgen", "缺少语句块"));
  }

  BlockFlow flow = BlockFlow::Continue;
  for (auto* item : ctx->blockItem()) {
    if (!item) continue;

    flow = VisitBlockItemResult(*item);
    if (flow == BlockFlow::Terminated) {
      break;
    }

    auto* cur = builder_.GetInsertBlock();
    std::cout << "[DEBUG] current insert block: "
          << (cur ? cur->GetName() : "<null>") << std::endl;
    if (cur && cur->HasTerminator()) {
      break;
    }
  }

  return flow;
}

// 类型安全的包装器
IRGenImpl::BlockFlow IRGenImpl::VisitBlockItemResult(
    SysYParser::BlockItemContext& item) {
  auto result = item.accept(this);// 调用 visitBlockItem，返回 std::any 包装的 BlockFlow
  return std::any_cast<BlockFlow>(result); // 解包为 BlockFlow
}
// 用于遍历块内项，返回是否继续访问后续项（如遇到 return/break/continue 则终止访问）
std::any IRGenImpl::visitBlockItem(SysYParser::BlockItemContext* ctx) {
  std::cout << "[DEBUG IRGEN] visitBlockItem: " << (ctx ? ctx->getText() : "<null>") << std::endl;
  if (!ctx) {
    throw std::runtime_error(FormatError("irgen", "缺少块内项"));
  }
  // 块内项可以是语句或声明，优先处理语句（如 return/break/continue 可能终止块内执行）
  if (ctx->stmt()) {
    return ctx->stmt()->accept(this); // 语句访问返回 BlockFlow，指示是否继续访问后续项
  }
  // 处理声明（如变量定义），继续访问后续项
  if (ctx->decl()) {
    ctx->decl()->accept(this);
    return BlockFlow::Continue; // 声明不会终止块内执行，继续访问后续项
  }
  throw std::runtime_error(FormatError("irgen", "暂不支持的块内项"));
}