"""
transpiler/restructurer.py — Reconstructs high-level structural IR from CFG using Region Analysis.
"""

from __future__ import annotations
from dataclasses import dataclass, field
from typing import List, Set, Dict, Optional, Union, Any, cast
from transpiler.middleend import cfg, ir, ssa
from transpiler.optimization import loop_detector
from transpiler.middleend.structural_analysis import PostDominatorAnalysis, LoopAnalyzer
from transpiler.backend import luau_ast as last



@dataclass
class Region:
    """Base class for all structured regions."""
    lineno: int = field(default=0, kw_only=True)

@dataclass
class BlockRegion(Region):
    """A region representing a single basic block."""
    block_id: int
    statements: List[ir.IRStatement] = field(default_factory=list)

@dataclass
class IfRegion(Region):
    """A region representing an if-then-else structure."""
    condition: ir.IRExpression
    then_region: Region
    else_region: Optional[Region] = None
    merge_block_id: Optional[int] = None

@dataclass
class LoopRegion(Region):
    """A region representing a cycle (while/for)."""
    header_id: int
    condition: Optional[ir.IRExpression]
    body: Region
    latch_id: int
    exit_id: Optional[int] = None   # primary exit block outside the loop
    extra_exits: Set[int] = field(default_factory=set)  # additional exit blocks (multi-exit)
    generic_for_info: Optional[Dict[str, Any]] = None

@dataclass
class SequenceRegion(Region):
    """A region representing a linear sequence of other regions."""
    regions: List[Region] = field(default_factory=list)

class CFGRestructurer:
    def __init__(self, target_cfg: cfg.CFG):
        self.cfg = target_cfg
        self.visited: Set[int] = set()

        # 1. Run Dominance and Post-Dominance
        self.dom = ssa.DominatorAnalysis(self.cfg)
        self.dom.run()
        self.pdom = PostDominatorAnalysis(self.cfg)
        self.pdom.run()

        # 2. Detect Loops
        self.loop_infos = loop_detector.detect_loops(self.cfg, self.dom)
        # Build block-to-loop mapping: block_id -> header_id of innermost loop
        self.block_to_loop: Dict[int, int] = {}
        for loop in self.loop_infos.values():
            for bid in loop.body_ids:
                bid_int = cast(int, bid)
                if bid_int in self.block_to_loop:
                    old_header = self.block_to_loop[bid_int]
                    if len(loop.body_ids) < len(self.loop_infos[old_header].body_ids):
                        self.block_to_loop[bid_int] = loop.header_id
                else:
                    self.block_to_loop[bid_int] = loop.header_id

    def reconstruct(self) -> ir.IRModule:
        root_region = self._build_region_from(self.cfg.entry_block.id)
        return ir.IRModule(body=self._emit_region(root_region))

    def _build_region_from(
        self,
        bid: int,
        stop_at: Optional[int] = None,
        loop_exit_ids: Optional[Set[int]] = None,
        loop_header_id: Optional[int] = None,
    ) -> Region:
        """Build a structured Region from basic block *bid*.

        Parameters
        ----------
        stop_at:
            Stop traversal when we reach this block (used for merge points,
            loop headers, etc.).
        loop_exit_ids:
            If the current block's terminator unconditionally jumps to one of
            these block IDs we are inside a loop body and this jump represents
            a `break`.  We emit a synthetic ``BlockRegion`` containing an
            ``IRBreakStatement``.
        loop_header_id:
            The header block of the enclosing loop.  A jump back to this block
            represents a `continue`.
        """
        if loop_exit_ids is None:
            loop_exit_ids = set()

        if stop_at is not None and bid == stop_at:
            return SequenceRegion()

        if bid in self.visited:
            return SequenceRegion()

        block = self.cfg.blocks[bid]
        # Is this a loop header?
        if bid in self.loop_infos and bid not in self.visited:
            return self._build_loop_region(bid, stop_at, loop_exit_ids, loop_header_id)

        self.visited.add(bid)

        # Collect non-terminator statements (but KEEP break/continue)
        stmt_list = [s for s in block.statements if not isinstance(s, (ir.IRJump, ir.IRBranch, ir.IRGenericForTerminator))]
        current_block_region = BlockRegion(block_id=bid, statements=stmt_list)

        terminator = block.terminator
        if isinstance(terminator, ir.IRJump):
            target = terminator.target_block_id

            # Jump back to loop header → continue (only if not already there)
            if loop_header_id is not None and target == loop_header_id:
                if not any(isinstance(s, ir.IRContinueStatement) for s in stmt_list):
                    return BlockRegion(
                        block_id=bid,
                        statements=stmt_list + [ir.IRContinueStatement()],
                    )
                return current_block_region

            # Jump to loop exit → break (only if not already there)
            if target in loop_exit_ids:
                if not any(isinstance(s, ir.IRBreakStatement) for s in stmt_list):
                    return BlockRegion(
                        block_id=bid,
                        statements=stmt_list + [ir.IRBreakStatement()],
                    )
                return current_block_region

            if target == stop_at or target in self.visited:
                return current_block_region
            
            # If the block already ends with break/continue, stop traversal
            if any(isinstance(s, (ir.IRBreakStatement, ir.IRContinueStatement)) for s in stmt_list):
                return current_block_region

            return SequenceRegion(regions=[
                current_block_region,
                self._build_region_from(target, stop_at, loop_exit_ids, loop_header_id),
            ])

        elif isinstance(terminator, ir.IRBranch):
            # One branch may jump to the loop exit (break) while the other
            # stays in the loop body — this is the multi-exit pattern.
            # One branch may jump to the loop exit (break) or header (continue)
            # Prioritize continue check
            true_is_cont = loop_header_id is not None and terminator.true_block_id == loop_header_id
            false_is_cont = loop_header_id is not None and terminator.false_block_id == loop_header_id
            
            true_is_exit = terminator.true_block_id in loop_exit_ids and not true_is_cont
            false_is_exit = terminator.false_block_id in loop_exit_ids and not false_is_cont

            # Use Post-Dominator for the merge point
            merge_id = self.pdom.idoms.get(bid)

            # Determine the "then" and "else" regions, injecting breaks where
            # an arm goes directly to a loop exit.
            then_region: Region
            else_region: Optional[Region]

            def is_simple_exit(target_id: int) -> bool:
                b = self.cfg.blocks[target_id]
                stmts = [s for s in b.statements if not isinstance(s, (ir.IRJump, ir.IRBranch, ir.IRGenericForTerminator))]
                return len(stmts) == 0

            if true_is_cont and is_simple_exit(terminator.true_block_id):
                then_region = BlockRegion(block_id=-1, statements=[ir.IRContinueStatement()])
                else_region = self._build_region_from(
                    terminator.false_block_id, stop_at=merge_id,
                    loop_exit_ids=loop_exit_ids, loop_header_id=loop_header_id,
                )
            elif false_is_cont and is_simple_exit(terminator.false_block_id):
                then_region = self._build_region_from(
                    terminator.true_block_id, stop_at=merge_id,
                    loop_exit_ids=loop_exit_ids, loop_header_id=loop_header_id,
                )
                else_region = BlockRegion(block_id=-1, statements=[ir.IRContinueStatement()])
            elif true_is_exit and is_simple_exit(terminator.true_block_id):
                then_region = BlockRegion(block_id=-1, statements=[ir.IRBreakStatement()])
                else_region = self._build_region_from(
                    terminator.false_block_id, stop_at=merge_id,
                    loop_exit_ids=loop_exit_ids, loop_header_id=loop_header_id,
                )
            elif false_is_exit and is_simple_exit(terminator.false_block_id):
                then_region = self._build_region_from(
                    terminator.true_block_id, stop_at=merge_id,
                    loop_exit_ids=loop_exit_ids, loop_header_id=loop_header_id,
                )
                else_region = BlockRegion(block_id=-1, statements=[ir.IRBreakStatement()])
            else:
                then_region = self._build_region_from(
                    terminator.true_block_id, stop_at=merge_id,
                    loop_exit_ids=loop_exit_ids, loop_header_id=loop_header_id,
                )
                else_region = self._build_region_from(
                    terminator.false_block_id, stop_at=merge_id,
                    loop_exit_ids=loop_exit_ids, loop_header_id=loop_header_id,
                )

            def is_terminal(reg: Region) -> bool:
                if isinstance(reg, BlockRegion):
                    return any(isinstance(s, (ir.IRBreakStatement, ir.IRContinueStatement, ir.IRReturnStatement)) for s in reg.statements)
                if isinstance(reg, SequenceRegion) and reg.regions:
                    return is_terminal(reg.regions[-1])
                return False

            if is_terminal(then_region) and else_region and not isinstance(else_region, SequenceRegion):
                # Guard Lifting: if a: return else: b -> if a: return end; b
                return SequenceRegion(regions=[
                    current_block_region,
                    IfRegion(
                        condition=terminator.condition,
                        then_region=then_region,
                        else_region=None,
                        merge_block_id=None,
                        lineno=terminator.lineno
                    ),
                    else_region
                ])

            if_region = IfRegion(
                condition=terminator.condition,
                then_region=then_region,
                else_region=else_region if terminator.false_block_id != merge_id else None,
                merge_block_id=merge_id,
                lineno=terminator.lineno
            )

            if merge_id is not None and merge_id != stop_at and (not loop_exit_ids or merge_id not in loop_exit_ids):
                # If both arms are terminal, don't recurse into merge
                def is_terminal(reg: Region) -> bool:
                    if isinstance(reg, BlockRegion):
                        return any(isinstance(s, (ir.IRBreakStatement, ir.IRContinueStatement, ir.IRReturnStatement)) for s in reg.statements)
                    if isinstance(reg, SequenceRegion) and reg.regions:
                        return is_terminal(reg.regions[-1])
                    return False
                
                if is_terminal(then_region) and (else_region is None or is_terminal(else_region)):
                     return SequenceRegion(regions=[current_block_region, if_region])

                seq = SequenceRegion(regions=[
                    current_block_region,
                    if_region,
                    self._build_region_from(merge_id, stop_at, loop_exit_ids, loop_header_id),
                ])
                return seq
            return SequenceRegion(regions=[current_block_region, if_region])

        return current_block_region

    def _build_loop_region(self, header_id: int, stop_at: Optional[int] = None, 
                           loop_exit_ids: Set[int] = set(), loop_header_id: Optional[int] = None) -> Region:
        info = self.loop_infos[header_id]
        analyzer = LoopAnalyzer(self.cfg, header_id, info.body_ids)
        analyzer.analyze()

        header_block = self.cfg.blocks[header_id]
        condition = None
        loop_exit_id = None
        generic_for_info = None

        # Collect ALL exit edges (jumps from a loop-body block to outside the loop)
        extra_exits: Set[int] = set()
        for body_bid in info.body_ids:
            body_block = self.cfg.blocks[body_bid]
            for succ in body_block.successors:
                if succ.id not in info.body_ids and succ.id != header_id:
                    extra_exits.add(succ.id)

        # Determine primary exit (the one the header branches to)
        if isinstance(header_block.terminator, ir.IRGenericForTerminator):
            term = header_block.terminator
            generic_for_info = {
                "vars": term.vars,
                "iterator": term.iterator
            }
            loop_exit_id = term.false_block_id
            condition = None
        elif isinstance(header_block.terminator, ir.IRBranch):
            term = header_block.terminator
            if condition is None and generic_for_info is None:
                if term.true_block_id not in info.body_ids:
                    loop_exit_id = term.true_block_id
                    condition = ir.IRUnaryOperation(operator="Not", operand=term.condition)
                else:
                    loop_exit_id = term.false_block_id
                    condition = term.condition
        
        # If we couldn't find an exit from the header (e.g. it's a Jump to the condition block),
        # use the single exit if one exists.
        if loop_exit_id is None and len(analyzer.exits) == 1:
            loop_exit_id = list(analyzer.exits)[0]

        # All exits from inside the loop (for break detection)
        # IMPORTANT: We MUST NOT include the current header_id in all_exits,
        # otherwise a continue to the current header might be seen as a break.
        all_exits: Set[int] = extra_exits.copy()
        if loop_exit_id is not None:
            all_exits.add(loop_exit_id)
        if header_id in all_exits:
            all_exits.remove(header_id)

        # Mark header as visited to avoid infinite recursion
        self.visited.add(header_id)

        # The header's own statements (non-terminators) belong at the top of the body
        header_stmts = [s for s in header_block.statements if not isinstance(s, (ir.IRJump, ir.IRBranch, ir.IRGenericForTerminator))]
        header_region = BlockRegion(block_id=header_id, statements=header_stmts)

        # Loop body starts at the successor that is INSIDE the loop
        body_start = None
        if isinstance(header_block.terminator, ir.IRBranch):
            term = header_block.terminator
            if term.true_block_id in info.body_ids:
                body_start = term.true_block_id
            elif term.false_block_id in info.body_ids:
                body_start = term.false_block_id
        elif header_block.successors:
            body_start = header_block.successors[0].id

        # The body is traversed with knowledge of all exits so break/continue
        # edges can be emitted correctly.
        body_region: Region = SequenceRegion()
        if body_start is not None and body_start != header_id:
            body_region = self._build_region_from(
                body_start,
                stop_at=header_id,
                loop_exit_ids=all_exits,
                loop_header_id=header_id,
            )

        l_region = LoopRegion(
            header_id=header_id,
            condition=condition,
            body=SequenceRegion(regions=[header_region, body_region]) if header_stmts else body_region,
            latch_id=list(analyzer.latches)[0] if analyzer.latches else header_id,
            exit_id=loop_exit_id if loop_exit_id is not None else stop_at,
            extra_exits=extra_exits,
            generic_for_info=generic_for_info if 'generic_for_info' in locals() else None,
            lineno=header_block.terminator.lineno if header_block.terminator else 0
        )

        # Continue emitting after the primary exit
        if loop_exit_id is not None and loop_exit_id != stop_at:
            after_region = self._build_region_from(loop_exit_id, stop_at, loop_exit_ids, loop_header_id)
            return SequenceRegion(regions=[l_region, after_region])
        return l_region

    def _emit_region(self, reg: Optional[Region]) -> List[ir.IRStatement]:
        if reg is None:
            return []
        if isinstance(reg, BlockRegion):
            return list(reg.statements)
        elif isinstance(reg, SequenceRegion):
            res: List[ir.IRStatement] = []
            for r in reg.regions:
                new_stmts = self._emit_region(r)
                res.extend(new_stmts)
            return res
        elif isinstance(reg, IfRegion):
            then_block = self._emit_region(reg.then_region)
            else_block = self._emit_region(reg.else_region)

            # Check for elseif pattern: flatten when the else has a single if
            # HOWEVER, for elite-tier guard merging, we avoid this if both are early returns,
            # allowing the peephole optimizer to merge them into a single 'if A or B then return end'.
            if len(else_block) == 1 and isinstance(else_block[0], ir.IRIfStatement):
                inner_if = else_block[0]
                is_outer_guard = (len(then_block) == 1 and isinstance(then_block[0], ir.IRReturnStatement))
                is_inner_guard = (len(inner_if.then_block) == 1 and isinstance(inner_if.then_block[0], ir.IRReturnStatement))
                
                if not (is_outer_guard and is_inner_guard):
                    return [ir.IRIfStatement(
                        condition=reg.condition,
                        then_block=then_block,
                        elif_blocks=[(inner_if.condition, inner_if.then_block)] + inner_if.elif_blocks,
                        else_block=inner_if.else_block,
                        lineno=reg.lineno
                    )]

            return [ir.IRIfStatement(
                condition=reg.condition,
                then_block=then_block,
                else_block=else_block,
                lineno=reg.lineno
            )]
        elif isinstance(reg, LoopRegion):
            body_stmts = self._emit_region(reg.body)
            if reg.generic_for_info:
                # Strip the synthetic MultiAssign added for SSA
                if body_stmts and isinstance(body_stmts[0], ir.IRMultiAssign):
                    body_stmts.pop(0)
                
                return [ir.IRGenericForStatement(
                    vars=reg.generic_for_info["vars"],
                    iterator=reg.generic_for_info["iterator"],
                    body=body_stmts,
                    lineno=reg.lineno
                )]
                
            return [ir.IRWhileStatement(
                condition=reg.condition if reg.condition else ir.IRLiteral(value=True),
                body=body_stmts,
                lineno=reg.lineno
            )]
        return []


def restructure_cfg(target_cfg: cfg.CFG) -> ir.IRModule:
    r = CFGRestructurer(target_cfg)
    return r.reconstruct()
