import sys
import os

sys.path.insert(0, os.path.join(os.path.dirname(__file__), ".."))

import ast
from transpiler.transformer import transform
from transpiler.backend.generator import generate
from transpiler.infrastructure.flags import CompilerFlags

def test_nested_loops_basic():
    source = """
def nested():
    count = 0
    for i in range(5):
        for j in range(5):
            if i + j > 5:
                break
            count += 1
    return count
"""
    tree = ast.parse(source)
    res = transform(tree)
    flags = CompilerFlags(no_runtime=True, fast=False)
    code_obj = generate(res, flags=flags)
    code = code_obj.code
    print("--- Nested Loops Basic ---")
    print(code)
    assert "break" in code
    # Ensure it's not 'break' for BOTH loops if only one has it
    assert code.count("break") == 1

def test_nested_loops_continue():
    source = """
def nested_continue():
    sum_val = 0
    for i in range(10):
        if i % 2 == 0:
            continue
        for j in range(3):
            if j == 1:
                continue
            sum_val += i + j
    return sum_val
"""
    tree = ast.parse(source)
    res = transform(tree)
    flags = CompilerFlags(no_runtime=True, fast=False)
    code_obj = generate(res, flags=flags)
    code = code_obj.code
    print("--- Nested Loops Continue ---")
    print(code)
    # The current restructurer might emit 'break' for top-level 'continue' if it's the only path
    # but we want to ensure 'continue' is actually used where possible.
    assert "continue" in code

def test_deeply_nested_loops():
    source = """
def deeply_nested():
    res = 0
    for i in range(3):
        for j in range(3):
            for k in range(3):
                if i == j == k:
                    break
                res += i * j * k
    return res
"""
    tree = ast.parse(source)
    res = transform(tree)
    flags = CompilerFlags(no_runtime=True, fast=False)
    code_obj = generate(res, flags=flags)
    code = code_obj.code
    print("--- Deeply Nested Loops ---")
    print(code)

if __name__ == "__main__":
    try:
        test_nested_loops_basic()
        test_nested_loops_continue()
        test_deeply_nested_loops()
        print("\n[PASS] All nested loop tests successfully built.")
    except Exception as e:
        print(f"\n[FAIL] Nested loop test failed: {e}")
        import traceback
        traceback.print_exc()
        sys.exit(1)
