from core.languages.python.cst_utils import parse_module_to_cst
from core.languages.python.testgen.postprocessing.range_modifier import modify_large_loops


def test_simple_large_range() -> None:
    test_source = """for i in range(2000):
    print(i)"""
    expected = """for i in range(1000):
    print(i)"""
    module = parse_module_to_cst(test_source)
    expected_module = parse_module_to_cst(expected)
    modified_module = modify_large_loops(module, threshold=1000)
    assert modified_module.code == expected_module.code


def test_power_operation() -> None:
    test_source = """for i in range(10**6):
    print(i)"""
    expected = """for i in range(1000):
    print(i)"""
    module = parse_module_to_cst(test_source)
    expected_module = parse_module_to_cst(expected)
    modified_module = modify_large_loops(module, threshold=1000)
    assert modified_module.code == expected_module.code


def test_complex_expression() -> None:
    test_source = """for i in range(2**3 * 500):
    print(i)"""
    expected = """for i in range(1000):
    print(i)"""
    module = parse_module_to_cst(test_source)
    expected_module = parse_module_to_cst(expected)
    modified_module = modify_large_loops(module, threshold=1000)
    assert modified_module.code == expected_module.code


def test_no_modification_needed() -> None:
    # Below threshold
    test_source = """for i in range(500):
    print(i)"""
    module = parse_module_to_cst(test_source)
    modified_module = modify_large_loops(module, threshold=1000)
    assert modified_module.deep_equals(module)


def test_nested_loops() -> None:
    test_source = """for i in range(2000):
    for j in range(3000):
        print(i, j)"""
    expected = """for i in range(1000):
    for j in range(1000):
        print(i, j)"""
    module = parse_module_to_cst(test_source)
    expected_module = parse_module_to_cst(expected)
    modified_module = modify_large_loops(module, threshold=1000)
    assert modified_module.code == expected_module.code


def test_custom_threshold() -> None:
    test_source = """for i in range(500):
    print(i)"""
    expected = """for i in range(100):
    print(i)"""
    module = parse_module_to_cst(test_source)
    expected_module = parse_module_to_cst(expected)
    modified_module = modify_large_loops(module, threshold=100)
    assert modified_module.code == expected_module.code


def test_list_comprehension() -> None:
    test_source = """x = [i for i in range(10000)]"""
    expected = """x = [i for i in range(1000)]"""
    module = parse_module_to_cst(test_source)
    expected_module = parse_module_to_cst(expected)
    modified_module = modify_large_loops(module, threshold=1000)
    assert modified_module.code == expected_module.code


def test_nested_list_comprehension() -> None:
    test_source = """x = [[i,j] for i in range(2000) for j in range(3000)]"""
    expected = """x = [[i,j] for i in range(1000) for j in range(1000)]"""
    module = parse_module_to_cst(test_source)
    expected_module = parse_module_to_cst(expected)
    modified_module = modify_large_loops(module, threshold=1000)
    assert modified_module.code == expected_module.code


def test_mixed_comprehensions() -> None:
    test_source = """
x = [i for i in range(2000)]
y = {i: i*2 for i in range(3000)}
z = (i*i for i in range(4000))"""
    expected = """
x = [i for i in range(1000)]
y = {i: i*2 for i in range(1000)}
z = (i*i for i in range(1000))"""
    module = parse_module_to_cst(test_source)
    expected_module = parse_module_to_cst(expected)
    modified_module = modify_large_loops(module, threshold=1000)
    assert modified_module.code == expected_module.code


# Test for multiple argument range() function
def test_two_argument_range() -> None:
    test_source = """
for i in range(500, 2500):
    print(i)"""
    expected = """
for i in range(500, 1500):
    print(i)"""
    module = parse_module_to_cst(test_source)
    expected_module = parse_module_to_cst(expected)
    modified_module = modify_large_loops(module, threshold=1000)
    assert modified_module.code == expected_module.code


def test_negative_range() -> None:
    test_source = """for i in range(0, -5000):
    print(i)"""
    expected = """for i in range(0, -1000):
    print(i)"""
    module = parse_module_to_cst(test_source)
    expected_module = parse_module_to_cst(expected)
    modified_module = modify_large_loops(module, threshold=1000)
    assert modified_module.code == expected_module.code


def test_range_with_step() -> None:
    test_source = """for i in range(0, 10000, 2):
    print(i)"""
    expected = """for i in range(0, 2000, 2):
    print(i)"""
    module = parse_module_to_cst(test_source)
    expected_module = parse_module_to_cst(expected)
    modified_module = modify_large_loops(module, threshold=1000)
    assert modified_module.code == expected_module.code


def test_negative_step() -> None:
    test_source = """for i in range(5000, 0, -1):
    print(i)"""
    expected = """for i in range(5000, 4000, -1):
    print(i)"""
    module = parse_module_to_cst(test_source)
    expected_module = parse_module_to_cst(expected)
    modified_module = modify_large_loops(module, threshold=1000)
    assert modified_module.code == expected_module.code


def test_negative_range_negative_step() -> None:
    test_source = """for i in range(-1000, -6000, -1):
    print(i)"""
    expected = """for i in range(-1000, -2000, -1):
    print(i)"""
    module = parse_module_to_cst(test_source)
    expected_module = parse_module_to_cst(expected)
    modified_module = modify_large_loops(module, threshold=1000)
    assert modified_module.code == expected_module.code


def test_complex_step_expression() -> None:
    test_source = """for i in range(0, 10000, 2**2):
    print(i)"""
    expected = """for i in range(0, 4000, 4):
    print(i)"""
    module = parse_module_to_cst(test_source)
    expected_module = parse_module_to_cst(expected)
    modified_module = modify_large_loops(module, threshold=1000)
    assert modified_module.code == expected_module.code


def test_mixed_complex_ranges() -> None:
    # Test with multiple evaluated ranges
    test_source = """x = [i for i in range(-100, 2000, 2)]
y = [i for i in range(1000, -1000, -1)]
z = [i for i in range(2**3 * 100, 2**3 * 500, 2**2)]"""
    expected = """x = [i for i in range(-100, 1900, 2)]
y = [i for i in range(1000, 0, -1)]
z = [i for i in range(2**3 * 100, 2**3 * 500, 2**2)]"""
    module = parse_module_to_cst(test_source)
    expected_module = parse_module_to_cst(expected)
    modified_module = modify_large_loops(module, threshold=1000)
    assert modified_module.code == expected_module.code


def test_no_modification_step_cases() -> None:
    test_source = """for i in range(0, 2000, 10): print(i)
for i in range(100, 0, -1): print(i)
for i in range(0, 100, -1): print(i)"""
    module = parse_module_to_cst(test_source)
    modified_module = modify_large_loops(module, threshold=1000)
    assert modified_module.deep_equals(module)


def test_nested_mixed_ranges() -> None:
    test_source = """for i in range(-1000, 1000):
    for j in range(1000, -1000, -2):
        for k in range(0, 3000, 3):
            print(i, j, k)"""
    expected = """for i in range(-1000, 0):
    for j in range(1000, -1000, -2):
        for k in range(0, 3000, 3):
            print(i, j, k)"""
    module = parse_module_to_cst(test_source)
    expected_module = parse_module_to_cst(expected)
    modified_module = modify_large_loops(module, threshold=1000)
    assert modified_module.code == expected_module.code


def test_comprehension_with_steps() -> None:
    test_source = """x = [i for i in range(0, 5000, 2)]
y = {i: i for i in range(1000, -1000, -1)}
z = (i for i in range(-5000, 5000, 5))"""
    expected = """x = [i for i in range(0, 2000, 2)]
y = {i: i for i in range(1000, 0, -1)}
z = (i for i in range(-5000, 0, 5))"""
    module = parse_module_to_cst(test_source)
    expected_module = parse_module_to_cst(expected)
    modified_module = modify_large_loops(module, threshold=1000)

    assert modified_module.code == expected_module.code


# Tests for floats
def test_simple_float_range() -> None:
    test_source = """for i in range(1500.7):
    print(i)"""
    expected = """for i in range(1000):
    print(i)"""
    module = parse_module_to_cst(test_source)
    expected_module = parse_module_to_cst(expected)
    modified_module = modify_large_loops(module, threshold=1000)
    assert modified_module.code == expected_module.code


def test_two_argument_float_range() -> None:
    test_source = """for i in range(100.5, 2100.7):
    print(i)"""
    expected = """for i in range(100, 1100):
    print(i)"""
    module = parse_module_to_cst(test_source)
    expected_module = parse_module_to_cst(expected)
    modified_module = modify_large_loops(module, threshold=1000)
    assert modified_module.code == expected_module.code


def test_three_argument_float_range() -> None:
    test_source = """for i in range(0.0, 5000.5, 2.5):
    print(i)"""
    expected = """for i in range(0, 2000, 2):
    print(i)"""
    module = parse_module_to_cst(test_source)
    expected_module = parse_module_to_cst(expected)
    modified_module = modify_large_loops(module, threshold=1000)
    assert modified_module.code == expected_module.code


def test_negative_float_range() -> None:
    test_source = """for i in range(-1500.8):
    print(i)"""
    expected = """for i in range(-1000):
    print(i)"""
    module = parse_module_to_cst(test_source)
    expected_module = parse_module_to_cst(expected)
    modified_module = modify_large_loops(module, threshold=1000)
    assert modified_module.code == expected_module.code


def test_mixed_float_int_range() -> None:
    test_source = """for i in range(100, 3000.5, 2):
    print(i)"""
    expected = """for i in range(100, 2100, 2):
    print(i)"""
    module = parse_module_to_cst(test_source)
    expected_module = parse_module_to_cst(expected)
    modified_module = modify_large_loops(module, threshold=1000)
    assert modified_module.code == expected_module.code


def test_float_expression_range() -> None:
    test_source = """for i in range(2.5 * 1000):
    print(i)"""
    expected = """for i in range(1000):
    print(i)"""
    module = parse_module_to_cst(test_source)
    expected_module = parse_module_to_cst(expected)
    modified_module = modify_large_loops(module, threshold=1000)
    assert modified_module.code == expected_module.code


# random.sample cases
def test_random_sample_with_range() -> None:
    test_source = """random.sample(range(100000), 100000)"""
    expected = """random.sample(range(1000), 1000)"""
    module = parse_module_to_cst(test_source)
    expected_module = parse_module_to_cst(expected)
    modified_module = modify_large_loops(module, threshold=1000)
    assert modified_module.code == expected_module.code


def test_random_sample_with_small_range() -> None:
    test_source = """random.sample(range(500), 100)"""
    expected = """random.sample(range(500), 100)"""
    module = parse_module_to_cst(test_source)
    expected_module = parse_module_to_cst(expected)
    modified_module = modify_large_loops(module, threshold=1000)
    assert modified_module.code == expected_module.code


def test_random_sample_with_large_k_small_range() -> None:
    test_source = """random.sample(range(100), 2000)"""
    expected = """random.sample(range(100), 1000)"""
    module = parse_module_to_cst(test_source)
    expected_module = parse_module_to_cst(expected)
    modified_module = modify_large_loops(module, threshold=1000)
    assert modified_module.code == expected_module.code


def test_random_sample_with_complex_range() -> None:
    test_source = """random.sample(range(2**10 * 10), 2000)"""
    expected = """random.sample(range(1000), 1000)"""
    module = parse_module_to_cst(test_source)
    expected_module = parse_module_to_cst(expected)
    modified_module = modify_large_loops(module, threshold=1000)
    assert modified_module.code == expected_module.code


def test_random_sample_in_assignment() -> None:
    test_source = """x = random.sample(range(5000), 2000)"""
    expected = """x = random.sample(range(1000), 1000)"""
    module = parse_module_to_cst(test_source)
    expected_module = parse_module_to_cst(expected)
    modified_module = modify_large_loops(module, threshold=1000)
    assert modified_module.code == expected_module.code


def test_multiple_random_samples() -> None:
    test_source = """
x = random.sample(range(5000), 2000)
y = random.sample(range(3000), 1500)
z = random.sample(range(500), 100)
"""
    expected = """
x = random.sample(range(1000), 1000)
y = random.sample(range(1000), 1000)
z = random.sample(range(500), 100)
"""
    module = parse_module_to_cst(test_source)
    expected_module = parse_module_to_cst(expected)
    modified_module = modify_large_loops(module, threshold=1000)
    assert modified_module.code == expected_module.code


def test_random_sample_with_negative_range() -> None:
    test_source = """random.sample(range(-5000, 5000), 2000)"""
    expected = """random.sample(range(-5000, -4000), 1000)"""
    module = parse_module_to_cst(test_source)
    expected_module = parse_module_to_cst(expected)
    modified_module = modify_large_loops(module, threshold=1000)
    assert modified_module.code == expected_module.code


def test_random_sample_with_step() -> None:
    test_source = """random.sample(range(0, 10000, 2), 3000)"""
    expected = """random.sample(range(0, 2000, 2), 1000)"""
    module = parse_module_to_cst(test_source)
    expected_module = parse_module_to_cst(expected)
    modified_module = modify_large_loops(module, threshold=1000)
    assert modified_module.code == expected_module.code
