"""
tests/test_canonicalizer_safety.py

Golden tests for the loop canonicalization safety gate.

Tests verify:
  - Safe loops (simple range pattern) ARE converted to ForStatement.
  - Unsafe loops (impure bound, computed step, body mutation of induction var)
    are LEFT as WhileStatement.
"""
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.optimization.canonicalizer import IRCanonicalizer
from transpiler.middleend import ir


def _transpile_to_ir(source: str) -> ir.IRModule:
    """Parse Python source → IR module (no optimizer, no SSA)."""
    tree = pyast.parse(source)
    return PythonToIRTransformer().transform(tree)


def _extract_first_stmt(module: ir.IRModule, expected_type):
    """Return the first top-level statement, asserting its type."""
    assert module.body, "IR body is empty"
    stmt = module.body[0]
    assert isinstance(stmt, expected_type), (
        f"Expected {expected_type.__name__}, got {type(stmt).__name__}"
    )
    return stmt


# ---------------------------------------------------------------------------
# SAFE cases — must be converted to IRForStatement
# ---------------------------------------------------------------------------

def test_safe_simple_range():
    """while i < N with literal N and i = i + 1 → for loop."""
    src = """\
i = 0
while i < 10:
    print(i)
    i = i + 1
"""
    module = _transpile_to_ir(src)
    canon = IRCanonicalizer()
    canon.canonicalize(module)
    # The init assignment (i=0) and the while should be folded into one for
    _extract_first_stmt(module, ir.IRForStatement)
    print("[PASS] test_safe_simple_range")


def test_safe_variable_bound():
    """while i < n with a pure variable bound → for loop (pure variable is safe)."""
    src = """\
n = 5
i = 0
while i < n:
    print(i)
    i = i + 1
"""
    module = _transpile_to_ir(src)
    canon = IRCanonicalizer()
    canon.canonicalize(module)
    # First stmt: n = 5; second: the for loop
    stmts = module.body
    assert len(stmts) >= 2
    assert isinstance(stmts[1], ir.IRForStatement), (
        f"Expected IRForStatement, got {type(stmts[1]).__name__}"
    )
    print("[PASS] test_safe_variable_bound")


def test_safe_lte_operator():
    """while i <= N with literal N → for loop."""
    src = """\
i = 0
while i <= 5:
    print(i)
    i = i + 1
"""
    module = _transpile_to_ir(src)
    IRCanonicalizer().canonicalize(module)
    _extract_first_stmt(module, ir.IRForStatement)
    print("[PASS] test_safe_lte_operator")


# ---------------------------------------------------------------------------
# UNSAFE cases — must remain as IRWhileStatement
# ---------------------------------------------------------------------------

def test_unsafe_impure_bound_call():
    """while i < len(arr) — bound is a function call → keep as while."""
    src = """\
arr = [1, 2, 3]
i = 0
while i < len(arr):
    print(arr[i])
    i = i + 1
"""
    module = _transpile_to_ir(src)
    canon = IRCanonicalizer()
    canon.canonicalize(module)
    # The while must be preserved; the assignments come before it
    while_stmts = [s for s in module.body if isinstance(s, ir.IRWhileStatement)]
    assert while_stmts, "Expected WhileStatement to be preserved (impure bound)"
    print("[PASS] test_unsafe_impure_bound_call")


def test_unsafe_computed_step():
    """i = i + f(i) — step is not a constant 1 → keep as while."""
    src = """\
i = 0
while i < 10:
    print(i)
    i = i + get_step()
"""
    module = _transpile_to_ir(src)
    canon = IRCanonicalizer()
    canon.canonicalize(module)
    while_stmts = [s for s in module.body if isinstance(s, ir.IRWhileStatement)]
    assert while_stmts, "Expected WhileStatement (computed step)"
    print("[PASS] test_unsafe_computed_step")


def test_unsafe_induction_var_mutated_in_body():
    """i is reassigned conditionally inside the body → keep as while."""
    src = """\
i = 0
while i < 10:
    if i == 5:
        i = 7
    i = i + 1
"""
    module = _transpile_to_ir(src)
    canon = IRCanonicalizer()
    canon.canonicalize(module)
    while_stmts = [s for s in module.body if isinstance(s, ir.IRWhileStatement)]
    assert while_stmts, "Expected WhileStatement (body mutates induction var)"
    print("[PASS] test_unsafe_induction_var_mutated_in_body")


def test_unsafe_non_zero_start_impure():
    """while i < func() with non-trivial bound → keep as while."""
    src = """\
i = 3
while i < compute_limit():
    print(i)
    i = i + 1
"""
    module = _transpile_to_ir(src)
    canon = IRCanonicalizer()
    canon.canonicalize(module)
    while_stmts = [s for s in module.body if isinstance(s, ir.IRWhileStatement)]
    assert while_stmts, "Expected WhileStatement (impure bound from call)"
    print("[PASS] test_unsafe_non_zero_start_impure")


if __name__ == "__main__":
    test_safe_simple_range()
    test_safe_variable_bound()
    test_safe_lte_operator()
    test_unsafe_impure_bound_call()
    test_unsafe_computed_step()
    test_unsafe_induction_var_mutated_in_body()
    test_unsafe_non_zero_start_impure()
    print("\n[ALL PASS] Canonicalization safety gate tests.")
