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

namespace ir {

namespace {

// Helper to perform DFS and gather all blocks in a natural loop
std::unordered_set<BasicBlock*> GetLoopBlocks(BasicBlock* B, BasicBlock* H) {
  std::unordered_set<BasicBlock*> loop;
  std::vector<BasicBlock*> worklist;
  
  loop.insert(H);
  if (B != H) {
    loop.insert(B);
    worklist.push_back(B);
  }
  
  while (!worklist.empty()) {
    auto* curr = worklist.back();
    worklist.pop_back();
    for (auto* pred : curr->GetPredecessors()) {
      if (loop.find(pred) == loop.end()) {
        loop.insert(pred);
        worklist.push_back(pred);
      }
    }
  }
  return loop;
}

// Check if an opcode is a pure hoisting candidate (pure arithmetic, comparisons, GEP, casts)
bool IsPureHoistingCandidate(Opcode op) {
  switch (op) {
    case Opcode::Add:
    case Opcode::Sub:
    case Opcode::Mul:
    case Opcode::ICmpEQ:
    case Opcode::ICmpNE:
    case Opcode::ICmpLT:
    case Opcode::ICmpGT:
    case Opcode::ICmpLE:
    case Opcode::ICmpGE:
    case Opcode::FAdd:
    case Opcode::FSub:
    case Opcode::FMul:
    case Opcode::FDiv:
    case Opcode::FCmpEQ:
    case Opcode::FCmpNE:
    case Opcode::FCmpLT:
    case Opcode::FCmpGT:
    case Opcode::FCmpLE:
    case Opcode::FCmpGE:
    case Opcode::ZExt:
    case Opcode::SIToFP:
    case Opcode::FPToSI:
    case Opcode::GEP:
      return true;
    default:
      return false;
  }
}

} // namespace

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

  // 1. Run DominatorTree Analysis
  DominatorTree dom_tree(func);
  dom_tree.Run();

  // 2. Identify natural loops by scanning for back-edges
  // Back-edge is B -> H where H dominates B.
  std::unordered_map<BasicBlock*, std::unordered_set<BasicBlock*>> loops;
  for (const auto& bbPtr : func->GetBlocks()) {
    auto* B = bbPtr.get();
    for (auto* H : B->GetSuccessors()) {
      if (dom_tree.Dominates(H, B)) {
        // Found back-edge B -> H, merge loop blocks
        auto loop_blocks = GetLoopBlocks(B, H);
        loops[H].insert(loop_blocks.begin(), loop_blocks.end());
      }
    }
  }

  // 3. Optimize each identified loop
  for (auto& pair : loops) {
    BasicBlock* H = pair.first;
    const auto& loop_blocks = pair.second;

    // A preheader is the single predecessor of H outside the loop
    BasicBlock* preheader = nullptr;
    int num_outside_preds = 0;
    for (auto* pred : H->GetPredecessors()) {
      if (loop_blocks.find(pred) == loop_blocks.end()) {
        preheader = pred;
        num_outside_preds++;
      }
    }

    // Hoist only if there is exactly one outside predecessor (which is the preheader)
    if (num_outside_preds != 1 || !preheader) {
      continue;
    }

    // Identify loop-invariant instructions
    std::unordered_set<Instruction*> invariant_insts;
    std::vector<Instruction*> invariant_order;
    bool local_changed = true;
    while (local_changed) {
      local_changed = false;

      for (auto* bb : loop_blocks) {
        for (const auto& instPtr : bb->GetInstructions()) {
          auto* inst = instPtr.get();
          
          if (invariant_insts.find(inst) != invariant_insts.end()) {
            continue; // Already identified
          }

          if (!IsPureHoistingCandidate(inst->GetOpcode())) {
            continue; // Cannot hoist impure instructions (load, store, call, branch)
          }

          // Check if all operands are loop-invariant
          bool all_ops_invariant = true;
          for (size_t i = 0; i < inst->GetNumOperands(); ++i) {
            auto* op = inst->GetOperand(i);
            
            // Constants are invariant
            if (dynamic_cast<ConstantValue*>(op)) {
              continue;
            }

            // Values defined outside the loop are invariant
            if (auto* op_inst = dynamic_cast<Instruction*>(op)) {
              if (loop_blocks.find(op_inst->GetParent()) == loop_blocks.end()) {
                continue;
              }
              // If defined inside the loop, must be already marked invariant
              if (invariant_insts.find(op_inst) != invariant_insts.end()) {
                continue;
              }
            } else {
              // Arguments and Globals are always defined outside the loop
              continue;
            }

            all_ops_invariant = false;
            break;
          }

          if (all_ops_invariant) {
            invariant_insts.insert(inst);
            invariant_order.push_back(inst);
            local_changed = true;
            changed = true;
          }
        }
      }
    }

    // Hoist the loop-invariant instructions into the preheader (in dependency order)
    for (auto* inst : invariant_order) {
      auto& source_insts = const_cast<std::vector<std::unique_ptr<Instruction>>&>(inst->GetParent()->GetInstructions());
      auto& preheader_insts = const_cast<std::vector<std::unique_ptr<Instruction>>&>(preheader->GetInstructions());

      std::unique_ptr<Instruction> moved_inst;
      for (auto it = source_insts.begin(); it != source_insts.end(); ++it) {
        if (it->get() == inst) {
          moved_inst = std::move(*it);
          source_insts.erase(it);
          break;
        }
      }

      if (moved_inst) {
        moved_inst->SetParent(preheader);
        // Insert right before the terminator branch instruction of the preheader block
        if (!preheader_insts.empty() && preheader->HasTerminator()) {
          auto* term = preheader_insts.back().get();
          preheader->InsertInstructionBefore(std::move(moved_inst), term);
        } else {
          preheader_insts.push_back(std::move(moved_inst));
        }
      }
    }
  }

  return changed;
}

} // namespace ir
