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

namespace ir {

// Helper to rebuild CFG predecessors and successors.
void RebuildCFG(Function* func) {
  for (auto& bbPtr : func->GetBlocks()) {
    bbPtr->ClearPredecessors();
    bbPtr->ClearSuccessors();
  }
  for (auto& bbPtr : func->GetBlocks()) {
    auto* bb = bbPtr.get();
    const auto& insts = bb->GetInstructions();
    if (insts.empty()) continue;
    auto* term = insts.back().get();
    if (auto* br = dynamic_cast<BranchInst*>(term)) {
      if (br->IsConditional()) {
        auto* t = br->GetIfTrue();
        auto* f = br->GetIfFalse();
        if (t) {
          bb->AddSuccessor(t);
          t->AddPredecessor(bb);
        }
        if (f) {
          bb->AddSuccessor(f);
          f->AddPredecessor(bb);
        }
      } else {
        auto* dest = br->GetDest();
        if (dest) {
          bb->AddSuccessor(dest);
          dest->AddPredecessor(bb);
        }
      }
    }
  }
}

static void PostOrderDFS(BasicBlock* bb, std::unordered_set<BasicBlock*>& visited,
                         std::vector<BasicBlock*>& post_order) {
  visited.insert(bb);
  for (auto* succ : bb->GetSuccessors()) {
    if (visited.find(succ) == visited.end()) {
      PostOrderDFS(succ, visited, post_order);
    }
  }
  post_order.push_back(bb);
}

DominatorTree::DominatorTree(Function* func) : func_(func) {}

void DominatorTree::Run() {
  RebuildCFG(func_);
  ComputeRPO();
  ComputeIdom();
  ComputeDomTree();
  ComputeDF();
}

void DominatorTree::ComputeRPO() {
  rpo_.clear();
  if (func_->GetBlocks().empty()) return;
  std::unordered_set<BasicBlock*> visited;
  std::vector<BasicBlock*> post_order;
  PostOrderDFS(func_->GetEntry(), visited, post_order);
  rpo_ = std::vector<BasicBlock*>(post_order.rbegin(), post_order.rend());
}

void DominatorTree::ComputeIdom() {
  idom_.clear();
  if (rpo_.empty()) return;

  BasicBlock* entry = rpo_.front();
  idom_[entry] = entry;

  std::unordered_map<BasicBlock*, int> rpo_index;
  for (size_t i = 0; i < rpo_.size(); ++i) {
    rpo_index[rpo_[i]] = i;
  }

  bool changed = true;
  while (changed) {
    changed = false;
    for (size_t i = 1; i < rpo_.size(); ++i) {
      BasicBlock* b = rpo_[i];
      BasicBlock* new_idom = nullptr;

      // Find first predecessor with a defined idom
      for (auto* pred : b->GetPredecessors()) {
        if (idom_.find(pred) != idom_.end()) {
          new_idom = pred;
          break;
        }
      }

      if (new_idom) {
        for (auto* pred : b->GetPredecessors()) {
          if (pred != new_idom && idom_.find(pred) != idom_.end()) {
            // Intersect
            auto* finger1 = pred;
            auto* finger2 = new_idom;
            while (finger1 != finger2) {
              while (rpo_index.at(finger1) > rpo_index.at(finger2)) {
                finger1 = idom_.at(finger1);
              }
              while (rpo_index.at(finger2) > rpo_index.at(finger1)) {
                finger2 = idom_.at(finger2);
              }
            }
            new_idom = finger1;
          }
        }

        if (idom_.find(b) == idom_.end() || idom_[b] != new_idom) {
          idom_[b] = new_idom;
          changed = true;
        }
      }
    }
  }
}

void DominatorTree::ComputeDomTree() {
  dom_tree_.clear();
  for (auto* b : rpo_) {
    dom_tree_[b] = {};
  }
  for (auto* b : rpo_) {
    if (b != rpo_.front()) {
      auto* parent = idom_[b];
      dom_tree_[parent].push_back(b);
    }
  }
}

void DominatorTree::ComputeDF() {
  df_.clear();
  for (auto* b : rpo_) {
    df_[b] = {};
  }
  for (auto* b : rpo_) {
    if (b->GetPredecessors().size() >= 2) {
      for (auto* pred : b->GetPredecessors()) {
        auto* runner = pred;
        auto* idom_b = idom_[b];
        while (runner != idom_b) {
          // If runner's df doesn't contain b already, add it
          auto& runner_df = df_[runner];
          if (std::find(runner_df.begin(), runner_df.end(), b) == runner_df.end()) {
            runner_df.push_back(b);
          }
          runner = idom_[runner];
        }
      }
    }
  }
}

BasicBlock* DominatorTree::GetIdom(BasicBlock* bb) const {
  auto it = idom_.find(bb);
  return it != idom_.end() ? it->second : nullptr;
}

const std::vector<BasicBlock*>& DominatorTree::GetDominatedBlocks(BasicBlock* bb) const {
  static const std::vector<BasicBlock*> empty;
  auto it = dom_tree_.find(bb);
  return it != dom_tree_.end() ? it->second : empty;
}

const std::vector<BasicBlock*>& DominatorTree::GetDominanceFrontier(BasicBlock* bb) const {
  static const std::vector<BasicBlock*> empty;
  auto it = df_.find(bb);
  return it != df_.end() ? it->second : empty;
}

bool DominatorTree::Dominates(BasicBlock* a, BasicBlock* b) const {
  if (a == b) return true;
  auto* runner = b;
  while (runner != rpo_.front()) {
    auto it = idom_.find(runner);
    if (it == idom_.end()) return false;
    runner = it->second;
    if (runner == a) return true;
  }
  return false;
}

} // namespace ir