#include "mir/MIR.h"
#include <unordered_map>
#include <unordered_set>
#include <vector>

namespace mir {

namespace {

int AlignTo(int value, int align) {
  return ((value + align - 1) / align) * align;
}

PhysReg NormalizeReg(PhysReg reg) {
  int r = static_cast<int>(reg);
  // Map 64-bit X0-X28 registers to 32-bit W0-W28 registers to handle aliasing
  if (r >= static_cast<int>(PhysReg::X0) && r <= static_cast<int>(PhysReg::X28)) {
    return static_cast<PhysReg>(r - static_cast<int>(PhysReg::X0) + static_cast<int>(PhysReg::W0));
  }
  return reg;
}

PhysReg MatchRegSize(PhysReg target, PhysReg src) {
  int t = static_cast<int>(target);
  int s = static_cast<int>(src);
  
  bool target_is_64 = (t >= static_cast<int>(PhysReg::X0) && t <= static_cast<int>(PhysReg::X28)) || 
                      t == static_cast<int>(PhysReg::X29) || 
                      t == static_cast<int>(PhysReg::X30) || 
                      t == static_cast<int>(PhysReg::SP);
  
  bool src_is_64 = (s >= static_cast<int>(PhysReg::X0) && s <= static_cast<int>(PhysReg::X28)) || 
                   s == static_cast<int>(PhysReg::X29) || 
                   s == static_cast<int>(PhysReg::X30) || 
                   s == static_cast<int>(PhysReg::SP);
  
  if (target_is_64 && !src_is_64) {
    if (s >= static_cast<int>(PhysReg::W0) && s <= static_cast<int>(PhysReg::W28)) {
      return static_cast<PhysReg>(s - static_cast<int>(PhysReg::W0) + static_cast<int>(PhysReg::X0));
    }
  } else if (!target_is_64 && src_is_64) {
    if (s >= static_cast<int>(PhysReg::X0) && s <= static_cast<int>(PhysReg::X28)) {
      return static_cast<PhysReg>(s - static_cast<int>(PhysReg::X0) + static_cast<int>(PhysReg::W0));
    }
  }
  return src;
}

bool IsFloatReg(PhysReg reg) {
  return reg >= PhysReg::S0 && reg <= PhysReg::S15;
}

bool SameReg(const Operand& lhs, const Operand& rhs) {
  return lhs.GetKind() == Operand::Kind::Reg &&
         rhs.GetKind() == Operand::Kind::Reg &&
         NormalizeReg(lhs.GetReg()) == NormalizeReg(rhs.GetReg());
}

bool IsZeroImm(const Operand& operand) {
  return operand.GetKind() == Operand::Kind::Imm && operand.GetImm() == 0;
}

std::vector<MachineInstr> SimplifyCompareToBranch(
    const std::vector<MachineInstr>& insts) {
  std::vector<MachineInstr> simplified;

  for (size_t i = 0; i < insts.size();) {
    if (i + 4 < insts.size() &&
        (insts[i].GetOpcode() == Opcode::CmpRR ||
         insts[i].GetOpcode() == Opcode::FCmpRR) &&
        insts[i + 1].GetOpcode() == Opcode::Cset &&
        insts[i + 2].GetOpcode() == Opcode::MovImm &&
        insts[i + 3].GetOpcode() == Opcode::CmpRR &&
        insts[i + 4].GetOpcode() == Opcode::BCond) {
      const auto& cset_ops = insts[i + 1].GetOperands();
      const auto& mov_ops = insts[i + 2].GetOperands();
      const auto& cmp2_ops = insts[i + 3].GetOperands();
      const auto& br_ops = insts[i + 4].GetOperands();

      if (mov_ops.size() == 2 && cmp2_ops.size() == 2 && br_ops.size() == 2 &&
          SameReg(cset_ops.at(0), cmp2_ops.at(0)) &&
          SameReg(mov_ops.at(0), cmp2_ops.at(1)) &&
          IsZeroImm(mov_ops.at(1)) &&
          br_ops.at(0).GetKind() == Operand::Kind::Cond &&
          br_ops.at(0).GetCondCode() == "ne") {
        simplified.push_back(insts[i]);
        simplified.emplace_back(Opcode::BCond,
                                std::vector<Operand>{
                                    Operand::Cond(cset_ops.at(1).GetCondCode()),
                                    br_ops.at(1)});
        i += 5;
        continue;
      }
    }

    simplified.push_back(insts[i]);
    i++;
  }

  return simplified;
}

void CompactFrameSlots(MachineFunction& function) {
  std::unordered_set<int> used_slots;
  for (const auto& block : function.GetBlocks()) {
    for (const auto& inst : block.GetInstructions()) {
      for (const auto& opnd : inst.GetOperands()) {
        if (opnd.GetKind() == Operand::Kind::FrameIndex) {
          used_slots.insert(opnd.GetFrameIndex());
        }
      }
    }
  }

  int cursor = 0;
  for (const auto& slot : function.GetFrameSlots()) {
    if (used_slots.find(slot.index) == used_slots.end()) {
      continue;
    }
    cursor += slot.size;
    function.GetFrameSlot(slot.index).offset = -cursor;
  }
  function.SetFrameSize(AlignTo(cursor, 16));
}

} // namespace

void RunPeephole(MachineFunction& function) {
  for (auto& block : function.GetBlocks()) {
    auto& insts = block.GetInstructions();
    insts = SimplifyCompareToBranch(insts);
    std::vector<MachineInstr> optimized;
    
    // Map from FrameIndex to the normalized physical register that currently holds its value
    std::unordered_map<int, PhysReg> slot_to_reg;

    for (const auto& inst : insts) {
      Opcode op = inst.GetOpcode();
      const auto& ops = inst.GetOperands();

      // 1. Handle register move elimination (e.g. mov w8, w8)
      if (op == Opcode::MovReg) {
        if (NormalizeReg(ops.at(0).GetReg()) == NormalizeReg(ops.at(1).GetReg())) {
          continue; // Delete redundant self-moves
        }
      }

      // 2. Handle redundant Load after Store
      if (op == Opcode::LoadStack) {
        int fi = ops.at(1).GetFrameIndex();
        auto it = slot_to_reg.find(fi);
        if (it != slot_to_reg.end()) {
          PhysReg source_reg = it->second;
          PhysReg dest_reg = NormalizeReg(ops.at(0).GetReg());
          if (source_reg == dest_reg) {
            // Loading the same register that already has the value - completely redundant!
            continue; 
          } else {
            // Replace LoadStack dest_reg, fi with MovReg dest_reg, matched_source
            PhysReg matched_source = MatchRegSize(ops.at(0).GetReg(), it->second);
            optimized.push_back(MachineInstr(Opcode::MovReg, {Operand::Reg(ops.at(0).GetReg()), Operand::Reg(matched_source)}));
            
            // Invalidate any other slots mapping to dest_reg because dest_reg is written
            std::vector<int> to_remove;
            for (const auto& pair : slot_to_reg) {
              if (NormalizeReg(pair.second) == dest_reg) {
                to_remove.push_back(pair.first);
              }
            }
            for (int key : to_remove) {
              slot_to_reg.erase(key);
            }
            
            // Add new mapping (normalized)
            slot_to_reg[fi] = dest_reg;
            continue;
          }
        }
      }

      // 3. Track and optimize stores
      if (op == Opcode::StoreStack) {
        PhysReg src = NormalizeReg(ops.at(0).GetReg());
        int fi = ops.at(1).GetFrameIndex();
        auto it = slot_to_reg.find(fi);
        if (it != slot_to_reg.end() && NormalizeReg(it->second) == src) {
          continue; // Delete redundant store
        }
        slot_to_reg[fi] = src;
      }

      // 4. Invalidate register mappings on writes
      bool writes_reg = false;
      PhysReg written_reg = PhysReg::W0; // dummy
      
      switch (op) {
        case Opcode::MovImm:
          if (!ops.empty() && ops.at(0).GetKind() == Operand::Kind::Reg) {
            writes_reg = true;
            written_reg = NormalizeReg(ops.at(0).GetReg());
            
            // Under the hood, MovImm to a float register implicitly writes to x8/w8
            if (IsFloatReg(ops.at(0).GetReg())) {
              PhysReg implicitly_written = NormalizeReg(PhysReg::X8);
              std::vector<int> to_remove;
              for (const auto& pair : slot_to_reg) {
                if (NormalizeReg(pair.second) == implicitly_written) {
                  to_remove.push_back(pair.first);
                }
              }
              for (int key : to_remove) {
                slot_to_reg.erase(key);
              }
            }
          }
          break;
        case Opcode::LoadStack:
        case Opcode::AddRR:
        case Opcode::SubRR:
        case Opcode::MulRR:
        case Opcode::SDivRR:
        case Opcode::MSubRRRR:
        case Opcode::FAddRRR:
        case Opcode::FSubRRR:
        case Opcode::FMulRRR:
        case Opcode::FDivRRR:
        case Opcode::Cset:
        case Opcode::MovReg:
        case Opcode::Adrp:
        case Opcode::AddRegImm:
        case Opcode::LslImm:
        case Opcode::LdrRegReg:
        case Opcode::SIToFP:
        case Opcode::FPToSI:
        case Opcode::ZExt:
          if (!ops.empty() && ops.at(0).GetKind() == Operand::Kind::Reg) {
            writes_reg = true;
            written_reg = NormalizeReg(ops.at(0).GetReg());
          }
          break;
        case Opcode::Call:
          // A function call destroys all temporary/scratch registers.
          slot_to_reg.clear();
          break;
        default:
          break;
      }

      if (writes_reg) {
        // Remove any slot mapping to this register
        std::vector<int> to_remove;
        for (const auto& pair : slot_to_reg) {
          if (NormalizeReg(pair.second) == written_reg) {
            to_remove.push_back(pair.first);
          }
        }
        for (int key : to_remove) {
          slot_to_reg.erase(key);
        }
      }

      optimized.push_back(inst);
    }
    
    insts = std::move(optimized);
  }
  
  // 5. Eliminate Dead Stack Slots (stores to slots that are never loaded or address-taken)
  // Count loads and address-taken operations
  std::unordered_map<int, int> load_count;
  std::unordered_map<int, int> address_taken_count;

  for (const auto& block : function.GetBlocks()) {
    for (const auto& inst : block.GetInstructions()) {
      Opcode op = inst.GetOpcode();
      const auto& ops = inst.GetOperands();
      
      for (const auto& opnd : ops) {
        if (opnd.GetKind() == Operand::Kind::FrameIndex) {
          int fi = opnd.GetFrameIndex();
          if (op == Opcode::LoadStack) {
            load_count[fi]++;
          } else if (op != Opcode::StoreStack) {
            address_taken_count[fi]++;
          }
        }
      }
    }
  }

  // Identify dead slots
  std::unordered_set<int> dead_slots;
  for (size_t i = 0; i < function.GetFrameSlots().size(); ++i) {
    int fi = static_cast<int>(i);
    if (load_count[fi] == 0 && address_taken_count[fi] == 0) {
      dead_slots.insert(fi);
    }
  }

  // Remove StoreStack to dead slots
  for (auto& block : function.GetBlocks()) {
    auto& insts = block.GetInstructions();
    std::vector<MachineInstr> optimized;
    for (const auto& inst : insts) {
      if (inst.GetOpcode() == Opcode::StoreStack) {
        int fi = inst.GetOperands().at(1).GetFrameIndex();
        if (dead_slots.find(fi) != dead_slots.end()) {
          continue; // Delete this store
        }
      }
      optimized.push_back(inst);
    }
    insts = std::move(optimized);
  }

  CompactFrameSlots(function);
}

} // namespace mir
