#include "ir/PassManager.h"
#include <iostream>
#include <unordered_map>
#include <vector>
#include <tuple>

namespace ir {

static bool IsEquivalent(Instruction* a, Instruction* b) {
  if (a->GetOpcode() != b->GetOpcode()) return false;
  if (a->GetNumOperands() != b->GetNumOperands()) return false;
  
  // Skip load, store, alloca, call, phi, branch, ret (since they have side-effects or special states)
  switch (a->GetOpcode()) {
    case Opcode::Add:
    case Opcode::Sub:
    case Opcode::Mul:
    case Opcode::Div:
    case Opcode::Mod:
    case Opcode::FAdd:
    case Opcode::FSub:
    case Opcode::FMul:
    case Opcode::FDiv:
    case Opcode::ICmpEQ:
    case Opcode::ICmpNE:
    case Opcode::ICmpLT:
    case Opcode::ICmpGT:
    case Opcode::ICmpLE:
    case Opcode::ICmpGE:
    case Opcode::FCmpEQ:
    case Opcode::FCmpNE:
    case Opcode::FCmpLT:
    case Opcode::FCmpGT:
    case Opcode::FCmpLE:
    case Opcode::FCmpGE:
    case Opcode::GEP:
    case Opcode::ZExt:
    case Opcode::SIToFP:
    case Opcode::FPToSI:
      break;
    default:
      return false; // Skip all other opcodes
  }

  // Compare all operands
  for (size_t i = 0; i < a->GetNumOperands(); ++i) {
    if (a->GetOperand(i) != b->GetOperand(i)) {
      return false;
    }
  }

  return true;
}

bool RunCSE(Function* func) {
  bool changed = false;

  for (const auto& bbPtr : func->GetBlocks()) {
    std::vector<Instruction*> seen_instructions;
    std::unordered_map<Value*, Instruction*> available_loads;
    std::vector<Instruction*> to_erase;

    for (const auto& instPtr : bbPtr->GetInstructions()) {
      auto* inst = instPtr.get();

      if (inst->GetOpcode() == Opcode::Load) {
        auto* load = static_cast<LoadInst*>(inst);
        auto it = available_loads.find(load->GetPtr());
        if (it != available_loads.end()) {
          inst->ReplaceAllUsesWith(it->second);
          to_erase.push_back(inst);
          changed = true;
          continue;
        }
        available_loads[load->GetPtr()] = inst;
        continue;
      }

      if (inst->GetOpcode() == Opcode::Store ||
          inst->GetOpcode() == Opcode::Call) {
        available_loads.clear();
      }

      Instruction* match = nullptr;
      for (auto* seen : seen_instructions) {
        if (IsEquivalent(inst, seen)) {
          match = seen;
          break;
        }
      }

      if (match) {
        inst->ReplaceAllUsesWith(match);
        to_erase.push_back(inst);
        changed = true;
      } else {
        seen_instructions.push_back(inst);
      }
    }

    for (auto* inst : to_erase) {
      bbPtr->EraseInstruction(inst);
    }
  }

  return changed;
}

} // namespace ir
