#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) {
        AddRuntimeFunctions();
      }

void IRGenImpl::AddRuntimeFunctions() {
  if (IRGen_DEBUG) std::cout << "[DEBUG IRGEN] 添加运行时库函数声明" << std::endl;
  
  // 输入函数（返回 int）
  module_.CreateFunction("getint", 
                         ir::Type::GetFunctionType(ir::Type::GetInt32Type(), {}));
  module_.CreateFunction("getch", 
                         ir::Type::GetFunctionType(ir::Type::GetInt32Type(), {}));
  
  // getarray(int* a, int n): int
  module_.CreateFunction("getarray", 
                         ir::Type::GetFunctionType(
                             ir::Type::GetInt32Type(),
                             {ir::Type::GetPtrInt32Type(), ir::Type::GetInt32Type()}));
  
  // 输出函数（返回 void）
  module_.CreateFunction("putint", 
                         ir::Type::GetFunctionType(
                             ir::Type::GetVoidType(),
                             {ir::Type::GetInt32Type()}));
  module_.CreateFunction("putch", 
                         ir::Type::GetFunctionType(
                             ir::Type::GetVoidType(),
                             {ir::Type::GetInt32Type()}));
  
  // putarray(int n, int* a): void
  module_.CreateFunction("putarray", 
                         ir::Type::GetFunctionType(
                             ir::Type::GetVoidType(),
                             {ir::Type::GetInt32Type(), ir::Type::GetPtrInt32Type()}));
  
  // 字符串输出（暂时用 int* 替代 char*）
  module_.CreateFunction("puts", 
                         ir::Type::GetFunctionType(
                             ir::Type::GetVoidType(),
                             {ir::Type::GetPtrInt32Type()}));
  
  // 时间测量函数（SysY 标准库）
  module_.CreateFunction("_sysy_starttime", 
                         ir::Type::GetFunctionType(
                             ir::Type::GetVoidType(),
                             {ir::Type::GetInt32Type()}));
  module_.CreateFunction("_sysy_stoptime", 
                         ir::Type::GetFunctionType(
                             ir::Type::GetVoidType(),
                             {ir::Type::GetInt32Type()}));
  
  // 简化版本
  module_.CreateFunction("starttime", 
                         ir::Type::GetFunctionType(ir::Type::GetVoidType(), {}));
  module_.CreateFunction("stoptime", 
                         ir::Type::GetFunctionType(ir::Type::GetVoidType(), {}));
  
  // 其他可能需要的函数
  module_.CreateFunction("read_map", 
                         ir::Type::GetFunctionType(ir::Type::GetInt32Type(), {}));
  
  // 浮点数
  module_.CreateFunction("float_eq", 
                         ir::Type::GetFunctionType(
                             ir::Type::GetInt32Type(),
                             {ir::Type::GetFloatType(), ir::Type::GetFloatType()}));
  
  // 内存操作函数
  module_.CreateFunction("memset", 
                         ir::Type::GetFunctionType(
                             ir::Type::GetPtrInt32Type(),
                             {ir::Type::GetPtrInt32Type(), 
                              ir::Type::GetInt32Type(), 
                              ir::Type::GetInt32Type()}));
  
  if (IRGen_DEBUG) std::cout << "[DEBUG IRGEN] 运行时库函数声明完成" << std::endl;
}

// 修正：没有 mainFuncDef，通过函数名找到 main
std::any IRGenImpl::visitCompUnit(SysYParser::CompUnitContext* ctx) {
  if (IRGen_DEBUG) 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) {
  if (IRGen_DEBUG) 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->Ident()) continue;
      std::string name = param->Ident()->getText();
      std::shared_ptr<ir::Type> param_ty;
      
      // 检查 bType 是否存在
      if (!param->bType()) {
        throw std::runtime_error(FormatError("irgen", "函数参数缺少类型: " + name));
      }
      
      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);
  
  // 调试输出
  if (IRGen_DEBUG) std::cerr << "[DEBUG] visitFuncDef: 创建函数 " << funcName 
            << "，返回类型: " << (ret_type->IsVoid() ? "void" : ret_type->IsFloat() ? "float" : "int")
            << "，参数数量: " << param_types.size() << std::endl;
  
  // 创建函数对象
  func_ = module_.CreateFunction(funcName, func_type);
  
  // 检查函数是否成功创建
  if (!func_) {
    std::cerr << "[ERROR] visitFuncDef: 创建函数失败，func_ 为 nullptr!" << std::endl;
    throw std::runtime_error(FormatError("irgen", "创建函数失败: " + funcName));
  }
  
  if (IRGen_DEBUG) std::cerr << "[DEBUG] visitFuncDef: 函数对象地址: " << (void*)func_ << std::endl;
  
  // 设置插入点
  auto* entry_block = func_->GetEntry();
  if (!entry_block) {
    std::cerr << "[ERROR] visitFuncDef: 函数入口基本块为空!" << std::endl;
    throw std::runtime_error(FormatError("irgen", "函数入口基本块为空: " + funcName));
  }
  
  builder_.SetInsertPoint(entry_block);
  storage_map_.clear();
  param_map_.clear();

  // 函数参数处理
  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;
      
      // 再次检查 bType
      if (!param->bType()) {
        throw std::runtime_error(FormatError("irgen", "函数参数缺少类型: " + name));
      }
      
      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();
        }
      }
      
      // 检查函数对象是否有效
      if (!func_) {
        std::cerr << "[ERROR] visitFuncDef: func_ 在添加参数时变为 nullptr!" << std::endl;
        throw std::runtime_error(FormatError("irgen", "函数对象无效"));
      }
      
      if (IRGen_DEBUG) 
 std::cerr << "[DEBUG] visitFuncDef: 为函数 " << funcName 
                << " 添加参数 " << name << "，类型: " 
                << (param_ty->IsInt32() ? "int32" : param_ty->IsFloat() ? "float" : 
                   param_ty->IsPtrInt32() ? "ptr_int32" : param_ty->IsPtrFloat() ? "ptr_float" : "other")
                << std::endl;
      
      // 创建参数并添加到函数
      auto arg = std::make_unique<ir::Argument>(param_ty, name);
      if (!arg) {
        throw std::runtime_error(FormatError("irgen", "创建参数失败: " + name));
      }
      
      auto* arg_ptr = arg.get();
      auto* added_arg = func_->AddArgument(std::move(arg));
      
      if (!added_arg) {
        std::cerr << "[ERROR] visitFuncDef: AddArgument 返回 nullptr!" << std::endl;
        throw std::runtime_error(FormatError("irgen", "添加参数失败: " + name));
      }
      
      // 为参数创建存储槽位
      ir::AllocaInst* slot = nullptr;
      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", "不支持的参数类型"));
      }
      
      if (!slot) {
        throw std::runtime_error(FormatError("irgen", "创建参数存储槽位失败: " + name));
      }
      
      builder_.CreateStore(added_arg, slot);
      param_map_[name] = slot;
      
      if (IRGen_DEBUG) 
 std::cerr << "[DEBUG] visitFuncDef: 参数 " << name << " 处理完成" << std::endl;
    }
  }

  // 生成函数体
  if (IRGen_DEBUG) std::cerr << "[DEBUG] visitFuncDef: 开始生成函数体" << std::endl;
  ctx->block()->accept(this);

  // 如果函数没有终止指令，添加默认返回
  if (!func_->GetEntry()->HasTerminator()) {
    if (IRGen_DEBUG)  std::cerr << "[DEBUG] visitFuncDef: 函数体没有终止指令，添加默认返回" << std::endl;
    auto retVal = builder_.CreateConstInt(0);
    builder_.CreateRet(retVal);
  }

  // 验证函数结构
  try {
    VerifyFunctionStructure(*func_);
  } catch (const std::exception& e) {
    std::cerr << "[ERROR] visitFuncDef: 验证函数结构失败: " << e.what() << std::endl;
    throw;
  }
  
  if (IRGen_DEBUG) std::cerr << "[DEBUG] visitFuncDef: 函数 " << funcName << " 生成完成" << std::endl;
  func_ = nullptr;
  return {};
}


std::any IRGenImpl::visitBlock(SysYParser::BlockContext* ctx) {
  if (IRGen_DEBUG) 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();
    if (IRGen_DEBUG) 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) {
  if (IRGen_DEBUG) 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", "暂不支持的块内项"));
}
