"""
tests/test_break_continue.py

Verifies that break and continue inside while-loops are correctly preserved
through the full compile pipeline (IR → CFG → SSA → Restructure → Canonicalize).
"""
import sys, os
sys.path.insert(0, os.path.join(os.path.dirname(__file__), ".."))

import ast as pyast
from transpiler.optimization.ir_transformer import PythonToIRTransformer
from transpiler.middleend.cfg_builder import CFGBuilder
from transpiler.backend.restructurer import restructure_cfg
from transpiler.optimization.canonicalizer import IRCanonicalizer
from transpiler.middleend import ir, ssa


def _compile(source: str) -> ir.IRModule:
    """Full pipeline: Python → IR → CFG → SSA → Restructure → Canonicalize."""
    tree = pyast.parse(source)
    ir_module = PythonToIRTransformer().transform(tree)
    ir_module = PythonToIRTransformer().transform(tree)
    target_cfg = CFGBuilder().build(ir_module)
    dom = ssa.DominatorAnalysis(target_cfg)
    dom.run()
    ssa.SSAConstructor(target_cfg, dom).run()
    
    restructured = restructure_cfg(target_cfg)
    
    IRCanonicalizer().canonicalize(restructured)
    return restructured


def _find_all(module: ir.IRModule, node_type) -> list:
    """Recursively collect all nodes of a given type from the IR."""
    results = []
    def walk(stmts):
        for s in stmts:
            if isinstance(s, node_type):
                results.append(s)
            for attr in ("body", "then_block", "else_block"):
                child = getattr(s, attr, None)
                if isinstance(child, list):
                    walk(child)
            for _cond, body in getattr(s, "elif_blocks", []):
                walk(body)
    walk(module.body)
    return results


# ---------------------------------------------------------------------------
# break tests
# ---------------------------------------------------------------------------

def test_simple_break():
    """A break inside a while-loop must survive the restructure pass."""
    src = """\
i = 0
while i < 10:
    if i == 5:
        break
    i = i + 1
"""
    module = _compile(src)
    breaks = _find_all(module, ir.IRBreakStatement)
    assert breaks, "Expected IRBreakStatement in output — break was lost"
    print("[PASS] test_simple_break")


def test_break_preserves_post_loop_code():
    """Code after the while loop must still be emitted after a break."""
    src = """\
i = 0
while i < 10:
    if i == 5:
        break
    i = i + 1
result = i
"""
    module = _compile(src)
    breaks = _find_all(module, ir.IRBreakStatement)
    assert breaks, "break was lost"
    # Check that there's at least one assignment to 'result' at top-level
    assignments = [s for s in module.body if isinstance(s, ir.IRAssignment)]
    names = [s.target.name if isinstance(s.target, ir.IRVariable) else "" for s in assignments]
    # After DeSSA 'result' might be versioned but original_name preserved
    found_result = any("result" in n for n in names)
    assert found_result, f"Post-loop assignment to 'result' missing; top-level names: {names}"
    print("[PASS] test_break_preserves_post_loop_code")


# ---------------------------------------------------------------------------
# continue tests
# ---------------------------------------------------------------------------

def test_simple_continue():
    """A continue inside a while-loop must survive the restructure pass."""
    src = """\
i = 0
total = 0
while i < 10:
    i = i + 1
    if i == 5:
        continue
    total = total + i
"""
    module = _compile(src)
    continues = _find_all(module, ir.IRContinueStatement)
    assert continues, "Expected IRContinueStatement in output — continue was lost"
    print("[PASS] test_simple_continue")


# ---------------------------------------------------------------------------
# multi-exit loop test
# ---------------------------------------------------------------------------

def test_multi_exit_loop():
    """A loop with two break paths must preserve at least one IRBreakStatement."""
    src = """\
i = 0
while i < 100:
    if i == 10:
        break
    if i == 20:
        break
    i = i + 1
"""
    module = _compile(src)
    breaks = _find_all(module, ir.IRBreakStatement)
    assert len(breaks) >= 1, f"Expected at least one IRBreakStatement, got {len(breaks)}"
    print(f"[PASS] test_multi_exit_loop (found {len(breaks)} break(s))")


if __name__ == "__main__":
    test_simple_break()
    test_break_preserves_post_loop_code()
    test_simple_continue()
    test_multi_exit_loop()
    print("\n[ALL PASS] Break / Continue tests.")
