"""
transpiler/loop_detector.py — Natural Loop Detection for CFG.
"""

from __future__ import annotations
from typing import Dict, Set, List, Optional, TYPE_CHECKING
from transpiler.middleend import cfg
if TYPE_CHECKING:
    from transpiler.middleend import ssa

class LoopInfo:
    def __init__(self, header_id: int):
        self.header_id = header_id
        self.body_ids: Set[int] = set()
        self.parent: Optional[LoopInfo] = None
        self.children: List[LoopInfo] = []

    def __repr__(self):
        return f"Loop(header={self.header_id}, body={sorted(list(self.body_ids))})"

class LoopDetector:
    def __init__(self, target_cfg: cfg.CFG, analysis: ssa.DominatorAnalysis):
        self.cfg = target_cfg
        self.analysis = analysis
        self.loops: Dict[int, LoopInfo] = {} # header_id -> LoopInfo

    def run(self) -> Dict[int, LoopInfo]:
        """Identifies natural loops in the CFG and returns a map of header_id -> LoopInfo."""
        # 1. Find back-edges and identify loops
        # A back-edge is an edge (u -> v) where v dominates u.
        for bid, block in self.cfg.blocks.items():
            for succ in block.successors:
                if succ.id in self.analysis.dominators.get(bid, set()):
                    self._add_to_loop(succ.id, bid)

        # 2. Reconstruct nesting structure
        self._nest_loops()
        return self.loops

    def _add_to_loop(self, header_id: int, back_edge_node_id: int):
        if header_id not in self.loops:
            self.loops[header_id] = LoopInfo(header_id)
        
        loop = self.loops[header_id]
        loop.body_ids.add(header_id)
        
        # Traverse backwards from the back-edge node to finding all nodes in the loop
        stack = [back_edge_node_id]
        while stack:
            n_id = stack.pop()
            if n_id not in loop.body_ids:
                loop.body_ids.add(n_id)
                block = self.cfg.blocks.get(n_id)
                if block:
                    for pred in block.predecessors:
                        stack.append(pred.id)

    def _nest_loops(self):
        """Builds a parent-child relationship tree for nested loops."""
        # Sort loops by body size (inner loops have smaller bodies)
        sorted_loops = sorted(self.loops.values(), key=lambda l: len(l.body_ids))
        for i, inner in enumerate(sorted_loops):
            for outer in sorted_loops[i+1:]:
                # Check if it's a strict subset to avoid self-nesting if header matches (rare)
                if inner.body_ids.issubset(outer.body_ids) and inner.header_id != outer.header_id:
                    inner.parent = outer
                    outer.children.append(inner)
                    break

def detect_loops(target_cfg: cfg.CFG, analysis: "ssa.DominatorAnalysis") -> Dict[int, LoopInfo]:
    detector = LoopDetector(target_cfg, analysis)
    return detector.run()
