// 极简 IR 定义：当前只支撑 i32 和加法，演示用。
// 可在此基础上扩展更多类型/指令
#pragma once

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

namespace ir {

class Type;
class ConstantInt;
class Instruction;
class BasicBlock;
class Function;

// IR 上下文：集中管理类型、常量等共享资源，便于复用与扩展。
class Context {
 public:
  Context() = default;
  ~Context();
  const std::shared_ptr<Type>& Void();
  const std::shared_ptr<Type>& Int32();
  const std::shared_ptr<Type>& PtrInt32();
  // 去重创建 i32 常量。
  ConstantInt* GetConstInt(int v);

  std::string NextTemp();

 private:
  std::shared_ptr<Type> void_;
  std::shared_ptr<Type> int32_;
  std::shared_ptr<Type> ptr_i32_;
  std::unordered_map<int, std::unique_ptr<ConstantInt>> const_ints_;
  int temp_index_ = -1;
};

class Type {
 public:
  enum class Kind { Void, Int32, PtrInt32 };
  explicit Type(Kind k);
  Kind GetKind() const;
  bool IsVoid() const;
  bool IsInt32() const;
  bool IsPtrInt32() const;

 private:
  Kind kind_;
};

class Value {
 public:
  Value(std::shared_ptr<Type> ty, std::string name);
  virtual ~Value() = default;
  const std::shared_ptr<Type>& GetType() const;
  const std::string& GetName() const;
  void SetName(std::string n);
  void AddUser(Instruction* user);
  const std::vector<Instruction*>& GetUsers() const;

 protected:
  std::shared_ptr<Type> type_;
  std::string name_;
  std::vector<Instruction*> users_;
};

class ConstantInt : public Value {
 public:
  ConstantInt(std::shared_ptr<Type> ty, int v);
  int GetValue() const { return value_; }

 private:
  int value_{};
};

// 后续还需要扩展更多指令类型。
enum class Opcode { Add, Sub, Mul, Alloca, Load, Store, Ret };

class Instruction : public Value {
 public:
  Instruction(Opcode op, std::shared_ptr<Type> ty, std::string name = "");
  Opcode GetOpcode() const;
  bool IsTerminator() const;
  BasicBlock* GetParent() const;
  void SetParent(BasicBlock* parent);

 private:
  Opcode opcode_;
  BasicBlock* parent_ = nullptr;
};

class BinaryInst : public Instruction {
 public:
  BinaryInst(Opcode op, std::shared_ptr<Type> ty, Value* lhs, Value* rhs,
             std::string name);
  Value* GetLhs() const;
  Value* GetRhs() const;

 private:
  Value* lhs_;
  Value* rhs_;
};

class ReturnInst : public Instruction {
 public:
  ReturnInst(std::shared_ptr<Type> void_ty, Value* val);
  Value* GetValue() const;

 private:
  Value* value_;
};

class AllocaInst : public Instruction {
 public:
  AllocaInst(std::shared_ptr<Type> ptr_ty, std::string name);
};

class LoadInst : public Instruction {
 public:
  LoadInst(std::shared_ptr<Type> val_ty, Value* ptr, std::string name);
  Value* GetPtr() const;

 private:
  Value* ptr_;
};

class StoreInst : public Instruction {
 public:
  StoreInst(std::shared_ptr<Type> void_ty, Value* val, Value* ptr);
  Value* GetValue() const;
  Value* GetPtr() const;

 private:
  Value* value_;
  Value* ptr_;
};

class BasicBlock {
 public:
  explicit BasicBlock(std::string name);
  const std::string& GetName() const;
  Function* GetParent() const;
  void SetParent(Function* parent);
  bool HasTerminator() const;
  const std::vector<std::unique_ptr<Instruction>>& GetInstructions() const;
  const std::vector<BasicBlock*>& GetPredecessors() const;
  const std::vector<BasicBlock*>& GetSuccessors() const;
  template <typename T, typename... Args>
  T* Append(Args&&... args) {
    if (HasTerminator()) {
      throw std::runtime_error("BasicBlock 已有 terminator，不能继续追加指令: " +
                               name_);
    }
    auto inst = std::make_unique<T>(std::forward<Args>(args)...);
    auto* ptr = inst.get();
    ptr->SetParent(this);
    instructions_.push_back(std::move(inst));
    return ptr;
  }

 private:
  std::string name_;
  Function* parent_ = nullptr;
  std::vector<std::unique_ptr<Instruction>> instructions_;
  std::vector<BasicBlock*> predecessors_;
  std::vector<BasicBlock*> successors_;
};

class Function : public Value {
 public:
  // 允许显式指定返回类型，便于后续扩展多种函数签名。
  Function(std::string name, std::shared_ptr<Type> ret_type);
  BasicBlock* CreateBlock(const std::string& name);
  BasicBlock* GetEntry();
  const BasicBlock* GetEntry() const;
  const std::vector<std::unique_ptr<BasicBlock>>& GetBlocks() const;

 private:
  BasicBlock* entry_ = nullptr;
  std::vector<std::unique_ptr<BasicBlock>> blocks_;
};

class Module {
 public:
  Module() = default;
  Context& GetContext();
  const Context& GetContext() const;
  // 创建函数时显式传入返回类型，便于在 IRGen 中根据语法树信息选择类型。
  Function* CreateFunction(const std::string& name,
                           std::shared_ptr<Type> ret_type);
  const std::vector<std::unique_ptr<Function>>& GetFunctions() const;

 private:
  Context context_;
  std::vector<std::unique_ptr<Function>> functions_;
};

class IRBuilder {
 public:
  IRBuilder(Context& ctx, BasicBlock* bb);
  void SetInsertPoint(BasicBlock* bb);
  BasicBlock* GetInsertBlock() const;

  // 构造常量、二元运算、返回指令的最小集合。
  ConstantInt* CreateConstInt(int v);
  BinaryInst* CreateBinary(Opcode op, Value* lhs, Value* rhs,
                           const std::string& name);
  BinaryInst* CreateAdd(Value* lhs, Value* rhs, const std::string& name);
  AllocaInst* CreateAllocaI32(const std::string& name);
  LoadInst* CreateLoad(Value* ptr, const std::string& name);
  StoreInst* CreateStore(Value* val, Value* ptr);
  ReturnInst* CreateRet(Value* v);

 private:
  Context& ctx_;
  BasicBlock* insert_block_;
};

class IRPrinter {
 public:
  void Print(const Module& module, std::ostream& os);
};

}  // namespace ir