#include "irgen/IRGen.h"

#include <stdexcept>

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

namespace {

std::shared_ptr<ir::Type> StorageType(std::shared_ptr<ir::Type> ty) {
  if (ty->IsInt32()) return ir::Type::GetPtrInt32Type();
  if (ty->IsFloat()) return ir::Type::GetPtrFloatType();
  return ty;
}

void VerifyFunctionStructure(const ir::Function& func) {
  for (const auto& bb : func.GetBlocks()) {
    if (!bb || !bb->HasTerminator()) {
      // If a block doesn't have a terminator, it might be an empty function or 
      // missing a return in a path. For SysY, we should at least have a default return for void.
      // But IRGen should have handled it.
      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),
      is_global_scope_(true) {}

std::any IRGenImpl::visitCompUnit(SysYParser::CompUnitContext* ctx) {
  if (!ctx) return {};
  
  is_global_scope_ = true;
  for (auto* decl : ctx->decl()) {
    decl->accept(this);
  }
  for (auto* funcDef : ctx->funcDef()) {
    funcDef->accept(this);
  }
  return {};
}

std::any IRGenImpl::visitFuncDef(SysYParser::FuncDefContext* ctx) {
  is_global_scope_ = false;
  
  std::shared_ptr<ir::Type> ret_ty;
  if (ctx->funcType()->INT()) {
    ret_ty = ir::Type::GetInt32Type();
  } else if (ctx->funcType()->FLOAT()) {
    ret_ty = ir::Type::GetFloatType();
  } else {
    ret_ty = ir::Type::GetVoidType();
  }

  std::string func_name = ctx->ID()->getText();
  
  std::vector<std::shared_ptr<ir::Type>> param_types;
  if (ctx->funcFParams()) {
    for (auto* fparam : ctx->funcFParams()->funcFParam()) {
      if (fparam->LBRACK().empty()) {
        param_types.push_back(fparam->btype()->INT() ? ir::Type::GetInt32Type() : ir::Type::GetFloatType());
      } else {
        // Array param is a pointer
        param_types.push_back(fparam->btype()->INT() ? ir::Type::GetPtrInt32Type() : ir::Type::GetPtrFloatType());
      }
    }
  }

  func_ = module_.CreateFunction(func_name, ret_ty, param_types);
  builder_.SetInsertPoint(func_->CreateBlock("entry"));

  // Handle parameters: alloca and store
  if (ctx->funcFParams()) {
    const auto& fparams = ctx->funcFParams()->funcFParam();
    const auto& args = func_->GetArguments();
    for (size_t i = 0; i < fparams.size(); ++i) {
      auto* fparam = fparams[i];
      auto* arg = args[i].get();
      auto* slot = builder_.CreateAlloca(StorageType(arg->GetType()), fparam->ID()->getText());
      builder_.CreateStore(arg, slot);
      storage_map_[fparam] = slot;
    }
  }

  ctx->blockStmt()->accept(this);

  // Default return for void functions if not terminated
  if (!builder_.GetInsertBlock()->HasTerminator()) {
    if (ret_ty->IsVoid()) {
      builder_.CreateRet(nullptr);
    } else if (ret_ty->IsInt32()) {
      builder_.CreateRet(builder_.CreateConstInt(0));
    } else if (ret_ty->IsFloat()) {
      builder_.CreateRet(builder_.CreateConstFloat(0.0f));
    }
  }

  VerifyFunctionStructure(*func_);
  is_global_scope_ = true;
  return {};
}

std::any IRGenImpl::visitFuncFParam(SysYParser::FuncFParamContext* ctx) {
  // We handle fparams in visitFuncDef directly.
  return {};
}
