#include "ir/PassManager.h"
#include <iostream>
#include <cmath>

namespace ir {

ConstantValue* FoldInstruction(Instruction* inst, Context& ctx) {
  if (inst->GetOpcode() == Opcode::ZExt) {
    auto* cast = static_cast<CastInst*>(inst);
    if (auto* ci = dynamic_cast<ConstantInt*>(cast->GetValue())) {
      return ctx.GetConstInt(ci->GetValue()); // ZExt is trivial on constant int
    }
  }

  if (inst->GetOpcode() == Opcode::SIToFP) {
    auto* cast = static_cast<CastInst*>(inst);
    if (auto* ci = dynamic_cast<ConstantInt*>(cast->GetValue())) {
      return ctx.GetConstFloat(static_cast<float>(ci->GetValue()));
    }
  }

  if (inst->GetOpcode() == Opcode::FPToSI) {
    auto* cast = static_cast<CastInst*>(inst);
    if (auto* cf = dynamic_cast<ConstantFloat*>(cast->GetValue())) {
      return ctx.GetConstInt(static_cast<int>(cf->GetValue()));
    }
  }

  // Binary operations
  if (auto* bin = dynamic_cast<BinaryInst*>(inst)) {
    auto* lhs = bin->GetLhs();
    auto* rhs = bin->GetRhs();

    auto* lhs_i = dynamic_cast<ConstantInt*>(lhs);
    auto* rhs_i = dynamic_cast<ConstantInt*>(rhs);
    auto* lhs_f = dynamic_cast<ConstantFloat*>(lhs);
    auto* rhs_f = dynamic_cast<ConstantFloat*>(rhs);

    if (lhs_i && rhs_i) {
      int l = lhs_i->GetValue();
      int r = rhs_i->GetValue();
      switch (bin->GetOpcode()) {
        case Opcode::Add: return ctx.GetConstInt(l + r);
        case Opcode::Sub: return ctx.GetConstInt(l - r);
        case Opcode::Mul: return ctx.GetConstInt(l * r);
        case Opcode::Div: return (r != 0) ? ctx.GetConstInt(l / r) : nullptr;
        case Opcode::Mod: return (r != 0) ? ctx.GetConstInt(l % r) : nullptr;
        case Opcode::ICmpEQ: return ctx.GetConstInt(l == r ? 1 : 0);
        case Opcode::ICmpNE: return ctx.GetConstInt(l != r ? 1 : 0);
        case Opcode::ICmpLT: return ctx.GetConstInt(l < r ? 1 : 0);
        case Opcode::ICmpGT: return ctx.GetConstInt(l > r ? 1 : 0);
        case Opcode::ICmpLE: return ctx.GetConstInt(l <= r ? 1 : 0);
        case Opcode::ICmpGE: return ctx.GetConstInt(l >= r ? 1 : 0);
        default: break;
      }
    }

    if (lhs_f && rhs_f) {
      float l = lhs_f->GetValue();
      float r = rhs_f->GetValue();
      switch (bin->GetOpcode()) {
        case Opcode::FAdd: return ctx.GetConstFloat(l + r);
        case Opcode::FSub: return ctx.GetConstFloat(l - r);
        case Opcode::FMul: return ctx.GetConstFloat(l * r);
        case Opcode::FDiv: return (r != 0.0f) ? ctx.GetConstFloat(l / r) : nullptr;
        case Opcode::FCmpEQ: return ctx.GetConstInt(l == r ? 1 : 0);
        case Opcode::FCmpNE: return ctx.GetConstInt(l != r ? 1 : 0);
        case Opcode::FCmpLT: return ctx.GetConstInt(l < r ? 1 : 0);
        case Opcode::FCmpGT: return ctx.GetConstInt(l > r ? 1 : 0);
        case Opcode::FCmpLE: return ctx.GetConstInt(l <= r ? 1 : 0);
        case Opcode::FCmpGE: return ctx.GetConstInt(l >= r ? 1 : 0);
        default: break;
      }
    }
  }

  return nullptr;
}

bool RunConstFold(Function* func, Context& ctx) {
  bool changed = false;
  std::vector<Instruction*> to_erase;

  for (const auto& bbPtr : func->GetBlocks()) {
    for (const auto& instPtr : bbPtr->GetInstructions()) {
      auto* inst = instPtr.get();
      if (inst->GetOpcode() == Opcode::Br || inst->GetOpcode() == Opcode::Ret || inst->GetOpcode() == Opcode::Phi) {
        continue;
      }
      if (auto* folded = FoldInstruction(inst, ctx)) {
        inst->ReplaceAllUsesWith(folded);
        to_erase.push_back(inst);
        changed = true;
      }
    }
  }

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

  return changed;
}

} // namespace ir
