#pragma once

#include <cstdint>
#include <iosfwd>
#include <memory>
#include <string>
#include <utility>
#include <vector>

namespace ir {
class Module;
}

namespace mir {

class MIRContext {
 public:
  MIRContext() = default;
};

MIRContext& DefaultContext();

enum class ValueType { Void, I1, I32, F32, Ptr };

enum class RegClass { GPR, FPR };

enum class CondCode { EQ, NE, LT, GT, LE, GE };

enum class StackObjectKind { Local, Spill, SavedGPR, SavedFPR };

enum class AddrBaseKind { None, FrameObject, Global, VReg };

enum class OperandKind { Invalid, VReg, Imm, Block, Symbol };

struct PhysReg {
  RegClass reg_class = RegClass::GPR;
  int index = -1;

  bool IsValid() const { return index >= 0; }
  bool operator==(const PhysReg& rhs) const {
    return reg_class == rhs.reg_class && index == rhs.index;
  }
};

bool IsGPR(ValueType type);
bool IsFPR(ValueType type);
int GetValueSize(ValueType type);
int GetValueAlign(ValueType type);
const char* GetPhysRegName(PhysReg reg, ValueType type);

class MachineOperand {
 public:
  MachineOperand() = default;
  static MachineOperand VReg(int reg);
  static MachineOperand Imm(std::int64_t value);
  static MachineOperand Block(std::string name);
  static MachineOperand Symbol(std::string name);

  OperandKind GetKind() const { return kind_; }
  int GetVReg() const { return vreg_; }
  std::int64_t GetImm() const { return imm_; }
  const std::string& GetText() const { return text_; }

 private:
  MachineOperand(OperandKind kind, int vreg, std::int64_t imm, std::string text);

  OperandKind kind_ = OperandKind::Invalid;
  int vreg_ = -1;
  std::int64_t imm_ = 0;
  std::string text_;
};

struct AddressExpr {
  AddrBaseKind base_kind = AddrBaseKind::None;
  int base_index = -1;
  std::string symbol;
  std::int64_t const_offset = 0;
  std::vector<std::pair<int, std::int64_t>> scaled_vregs;
};

struct StackObject {
  int index = -1;
  StackObjectKind kind = StackObjectKind::Local;
  int size = 0;
  int align = 1;
  int offset = 0;
  std::string name;
};

struct VRegInfo {
  int id = -1;
  ValueType type = ValueType::Void;
};

struct Allocation {
  enum class Kind { Unassigned, PhysReg, Spill };

  Kind kind = Kind::Unassigned;
  PhysReg phys;
  int stack_object = -1;
};

class MachineInstr {
 public:
  enum class Opcode {
    Arg,
    Copy,
    Load,
    Store,
    Lea,
    Add,
    Sub,
    Mul,
    Div,
    Rem,
    And,
    Or,
    Xor,
    Shl,
    AShr,
    LShr,
    FAdd,
    FSub,
    FMul,
    FDiv,
    FNeg,
    ICmp,
    FCmp,
    ZExt,
    ItoF,
    FtoI,
    Br,
    CondBr,
    Call,
    Ret,
    Memset,
    Unreachable,
  };

  explicit MachineInstr(Opcode opcode,
                        std::vector<MachineOperand> operands = {});

  Opcode GetOpcode() const { return opcode_; }
  const std::vector<MachineOperand>& GetOperands() const { return operands_; }
  std::vector<MachineOperand>& GetOperands() { return operands_; }

  void SetCondCode(CondCode code) { cond_code_ = code; }
  CondCode GetCondCode() const { return cond_code_; }

  void SetAddress(AddressExpr address) {
    address_ = std::move(address);
    has_address_ = true;
  }
  bool HasAddress() const { return has_address_; }
  const AddressExpr& GetAddress() const { return address_; }
  AddressExpr& GetAddress() { return address_; }

  void SetCallInfo(std::string callee, std::vector<ValueType> arg_types,
                   ValueType return_type) {
    callee_ = std::move(callee);
    call_arg_types_ = std::move(arg_types);
    call_return_type_ = return_type;
  }
  const std::string& GetCallee() const { return callee_; }
  const std::vector<ValueType>& GetCallArgTypes() const { return call_arg_types_; }
  ValueType GetCallReturnType() const { return call_return_type_; }

  void SetValueType(ValueType type) { value_type_ = type; }
  ValueType GetValueType() const { return value_type_; }

  bool IsTerminator() const;
  std::vector<int> GetDefs() const;
  std::vector<int> GetUses() const;

 private:
  Opcode opcode_;
  std::vector<MachineOperand> operands_;
  CondCode cond_code_ = CondCode::EQ;
  AddressExpr address_;
  bool has_address_ = false;
  std::string callee_;
  std::vector<ValueType> call_arg_types_;
  ValueType call_return_type_ = ValueType::Void;
  ValueType value_type_ = ValueType::Void;
};

class MachineBasicBlock {
 public:
  explicit MachineBasicBlock(std::string name);

  const std::string& GetName() const { return name_; }
  std::vector<MachineInstr>& GetInstructions() { return instructions_; }
  const std::vector<MachineInstr>& GetInstructions() const { return instructions_; }

  MachineInstr& Append(MachineInstr::Opcode opcode,
                       std::vector<MachineOperand> operands = {});
  MachineInstr& Append(MachineInstr instr);

 private:
  std::string name_;
  std::vector<MachineInstr> instructions_;
};

class MachineFunction {
 public:
  MachineFunction(std::string name, ValueType return_type,
                  std::vector<ValueType> param_types);

  const std::string& GetName() const { return name_; }
  ValueType GetReturnType() const { return return_type_; }
  const std::vector<ValueType>& GetParamTypes() const { return param_types_; }

  MachineBasicBlock* CreateBlock(const std::string& name);
  std::vector<std::unique_ptr<MachineBasicBlock>>& GetBlocks() { return blocks_; }
  const std::vector<std::unique_ptr<MachineBasicBlock>>& GetBlocks() const {
    return blocks_;
  }

  int NewVReg(ValueType type);
  const VRegInfo& GetVRegInfo(int id) const;
  VRegInfo& GetVRegInfo(int id);
  const std::vector<VRegInfo>& GetVRegs() const { return vregs_; }

  int CreateStackObject(int size, int align, StackObjectKind kind,
                        const std::string& name = "");
  StackObject& GetStackObject(int index);
  const StackObject& GetStackObject(int index) const;
  const std::vector<StackObject>& GetStackObjects() const { return stack_objects_; }

  void SetAllocation(int vreg, Allocation allocation);
  const Allocation& GetAllocation(int vreg) const;
  Allocation& GetAllocation(int vreg);

  void AddUsedCalleeSavedGPR(int reg_index);
  void AddUsedCalleeSavedFPR(int reg_index);
  const std::vector<int>& GetUsedCalleeSavedGPRs() const {
    return used_callee_saved_gprs_;
  }
  const std::vector<int>& GetUsedCalleeSavedFPRs() const {
    return used_callee_saved_fprs_;
  }

  void SetFrameSize(int size) { frame_size_ = size; }
  int GetFrameSize() const { return frame_size_; }

  void SetMaxOutgoingArgBytes(int bytes) { max_outgoing_arg_bytes_ = bytes; }
  int GetMaxOutgoingArgBytes() const { return max_outgoing_arg_bytes_; }

 private:
  std::string name_;
  ValueType return_type_ = ValueType::Void;
  std::vector<ValueType> param_types_;
  std::vector<std::unique_ptr<MachineBasicBlock>> blocks_;
  std::vector<VRegInfo> vregs_;
  std::vector<StackObject> stack_objects_;
  std::vector<Allocation> allocations_;
  std::vector<int> used_callee_saved_gprs_;
  std::vector<int> used_callee_saved_fprs_;
  int frame_size_ = 0;
  int max_outgoing_arg_bytes_ = 0;
};

class MachineModule {
 public:
  explicit MachineModule(const ir::Module& source) : source_(&source) {}

  const ir::Module& GetSourceModule() const { return *source_; }
  std::vector<std::unique_ptr<MachineFunction>>& GetFunctions() { return functions_; }
  const std::vector<std::unique_ptr<MachineFunction>>& GetFunctions() const {
    return functions_;
  }

  MachineFunction* AddFunction(std::unique_ptr<MachineFunction> function);

 private:
  const ir::Module* source_ = nullptr;
  std::vector<std::unique_ptr<MachineFunction>> functions_;
};

std::unique_ptr<MachineModule> LowerToMIR(const ir::Module& module);
void RunRegAlloc(MachineModule& module);
void RunFrameLowering(MachineModule& module);
void PrintAsm(const MachineModule& module, std::ostream& os);

}  // namespace mir