#include "irgen/IRGen.h"

#include <memory>
#include <stdexcept>
#include <string>
#include <unordered_map>
#include <utility>
#include <vector>

#include "ast/AstNodes.h"
#include "ir/IR.h"

namespace {

class IRGenImpl {
 public:
  explicit IRGenImpl(ir::Module& module)
      : module_(module),
        func_(module_.CreateFunction("main")),
        builder_(func_->entry()) {}

  void Gen(const ast::CompUnit& ast) {
    if (!ast.func || !ast.func->body) {
      throw std::runtime_error("AST 不完整：缺少 main 定义");
    }
    GenBlock(*ast.func->body);
  }

  std::unique_ptr<ir::Module> TakeModule() {
    return std::make_unique<ir::Module>(std::move(module_));
  }

 private:
  void GenBlock(const ast::Block& block) {
    for (const auto& decl : block.varDecls) {
      ir::Value* init = nullptr;
      if (decl->init) {
        init = GenExpr(*decl->init);
      } else {
        const_pool_.push_back(std::make_unique<ir::ConstantInt>(0));
        init = const_pool_.back().get();
      }
      locals_[decl->name] = init;
    }
    for (const auto& stmt : block.stmts) {
      if (auto ret = dynamic_cast<ast::ReturnStmt*>(stmt.get())) {
        ir::Value* v = GenExpr(*ret->value);
        builder_.CreateRet(v);
      }
    }
  }

  ir::Value* GenExpr(const ast::Expr& expr) {
    if (auto num = dynamic_cast<const ast::NumberExpr*>(&expr)) {
      const_pool_.push_back(std::make_unique<ir::ConstantInt>(num->value));
      return const_pool_.back().get();
    }
    if (auto var = dynamic_cast<const ast::VarExpr*>(&expr)) {
      auto it = locals_.find(var->name);
      if (it == locals_.end()) {
        throw std::runtime_error("变量未找到: " + var->name);
      }
      return it->second;
    }
    if (auto bin = dynamic_cast<const ast::BinaryExpr*>(&expr)) {
      auto* lhs = GenExpr(*bin->lhs);
      auto* rhs = GenExpr(*bin->rhs);
      std::string name = "%t" + std::to_string(temp_index_++);
      if (bin->op == ast::BinaryOp::Add) {
        return builder_.CreateBinary(ir::Opcode::Add, lhs, rhs, name);
      }
      if (bin->op == ast::BinaryOp::Sub) {
        // 当前子集只需要加法，减法复用 add 但保留分支，便于扩展
        return builder_.CreateBinary(ir::Opcode::Add, lhs, rhs, name);
      }
    }
    throw std::runtime_error("不支持的表达式类型");
  }

  ir::Module& module_;
  ir::Function* func_;
  ir::IRBuilder builder_;
  std::unordered_map<std::string, ir::Value*> locals_;
  std::vector<std::unique_ptr<ir::ConstantInt>> const_pool_;
  int temp_index_ = 0;
};

}  // namespace

std::unique_ptr<ir::Module> GenerateIR(const ast::CompUnit& ast) {
  auto module = std::make_unique<ir::Module>();
  IRGenImpl gen(*module);
  gen.Gen(ast);
  return module;
}