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

namespace ir {

// Predeclaration of CFG rebuild helper
void RebuildCFG(Function* func);

bool RunCFGSimplify(Function* func) {
  bool changed = false;
  bool local_changed = true;

  while (local_changed) {
    local_changed = false;
    RebuildCFG(func);

    // 1. Remove unreachable basic blocks
    BasicBlock* entry = func->GetEntry();
    std::unordered_set<BasicBlock*> reachable;
    std::queue<BasicBlock*> worklist;

    reachable.insert(entry);
    worklist.push(entry);
    while (!worklist.empty()) {
      auto* curr = worklist.front();
      worklist.pop();
      for (auto* succ : curr->GetSuccessors()) {
        if (reachable.find(succ) == reachable.end()) {
          reachable.insert(succ);
          worklist.push(succ);
        }
      }
    }

    std::vector<BasicBlock*> unreachable_blocks;
    for (const auto& bbPtr : func->GetBlocks()) {
      if (reachable.find(bbPtr.get()) == reachable.end()) {
        unreachable_blocks.push_back(bbPtr.get());
      }
    }

    if (!unreachable_blocks.empty()) {
      changed = true;
      local_changed = true;
      for (auto* bb : unreachable_blocks) {
        // Remove bb from predecessors of its successors, and clean up successor phi nodes
        for (auto* succ : bb->GetSuccessors()) {
          for (const auto& instPtr : succ->GetInstructions()) {
            if (instPtr->GetOpcode() == Opcode::Phi) {
              auto* phi = static_cast<PhiInst*>(instPtr.get());
              phi->RemoveIncomingBlock(bb);
            }
          }
        }

        // Remove from func's blocks
        auto& blocks = const_cast<std::vector<std::unique_ptr<BasicBlock>>&>(func->GetBlocks());
        blocks.erase(std::remove_if(blocks.begin(), blocks.end(),
                                    [&](const std::unique_ptr<BasicBlock>& b) {
                                      return b.get() == bb;
                                    }),
                     blocks.end());
      }
      continue; // Restart simplification loop safely
    }

    // 2. Merge basic block B with successor S if S has only one predecessor B
    for (const auto& bbPtr : func->GetBlocks()) {
      auto* b = bbPtr.get();
      if (b->GetSuccessors().size() == 1) {
        auto* s = b->GetSuccessors().front();
        if (s != entry && s->GetPredecessors().size() == 1) {
          changed = true;
          local_changed = true;

          // Replace all uses of block S as label with block B
          s->ReplaceAllUsesWith(b);

          // Erase B's terminator (the BranchInst to S)
          auto* b_term = b->GetInstructions().back().get();
          b->EraseInstruction(b_term);

          // For any PhiInst in S: it has exactly 1 incoming value from B.
          // Replace all uses of the PhiInst with its single incoming value.
          std::vector<Instruction*> phi_to_remove;
          for (const auto& instPtr : s->GetInstructions()) {
            if (instPtr->GetOpcode() == Opcode::Phi) {
              auto* phi = static_cast<PhiInst*>(instPtr.get());
              if (phi->GetNumIncoming() > 0) {
                phi->ReplaceAllUsesWith(phi->GetIncomingValue(0));
              }
              phi_to_remove.push_back(phi);
            }
          }

          // Move instructions from S to B
          auto& s_insts = const_cast<std::vector<std::unique_ptr<Instruction>>&>(s->GetInstructions());
          for (auto& instPtr : s_insts) {
            if (std::find(phi_to_remove.begin(), phi_to_remove.end(), instPtr.get()) == phi_to_remove.end()) {
              instPtr->SetParent(b);
              const_cast<std::vector<std::unique_ptr<Instruction>>&>(b->GetInstructions()).push_back(std::move(instPtr));
            }
          }

          // Clear S's instructions to prevent any dangling or double frees
          s_insts.clear();

          // Erase S from func's blocks list
          auto& blocks = const_cast<std::vector<std::unique_ptr<BasicBlock>>&>(func->GetBlocks());
          blocks.erase(std::remove_if(blocks.begin(), blocks.end(),
                                      [&](const std::unique_ptr<BasicBlock>& b) {
                                        return b.get() == s;
                                      }),
                       blocks.end());
          break; // Break to restart loop safely
        }
      }
    }
  }

  return changed;
}

} // namespace ir
