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

namespace ir {

// Declare FoldInstruction from ConstFold.cpp
ConstantValue* FoldInstruction(Instruction* inst, Context& ctx);

bool RunConstProp(Function* func, Context& ctx) {
  bool changed = false;
  bool local_changed = true;

  while (local_changed) {
    local_changed = false;
    std::vector<Instruction*> to_erase;

    // 1. Fold instructions
    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);
          local_changed = true;
          changed = true;
        }
      }
    }

    // Erase the folded instructions
    for (auto* inst : to_erase) {
      inst->GetParent()->EraseInstruction(inst);
    }

    // 2. Simplify conditional branches
    for (const auto& bbPtr : func->GetBlocks()) {
      auto* bb = bbPtr.get();
      const auto& insts = bb->GetInstructions();
      if (insts.empty()) continue;
      auto* term = insts.back().get();
      if (term->GetOpcode() == Opcode::Br) {
        auto* br = static_cast<BranchInst*>(term);
        if (br->IsConditional()) {
          if (auto* cond_const = dynamic_cast<ConstantInt*>(br->GetCondition())) {
            BasicBlock* target = (cond_const->GetValue() != 0) ? br->GetIfTrue() : br->GetIfFalse();
            BasicBlock* dead_target = (cond_const->GetValue() != 0) ? br->GetIfFalse() : br->GetIfTrue();

            if (dead_target != target) {
              for (const auto& instPtr : dead_target->GetInstructions()) {
                if (instPtr->GetOpcode() == Opcode::Phi) {
                  auto* phi = static_cast<PhiInst*>(instPtr.get());
                  phi->RemoveIncomingBlock(bb);
                }
              }
            }

            bb->EraseInstruction(br);
            bb->Append<BranchInst>(target);
            local_changed = true;
            changed = true;
            break; // Restart loop to handle CFG shifts safely
          }
        }
      }
    }
  }

  return changed;
}

} // namespace ir
