import libcst as cst

from core.languages.python.testgen.postprocessing.remove_unused_definitions import (
    remove_unused_definitions_from_pytest_file,
)


def test_basic_removal() -> None:
    """Test that unused functions and variables are removed."""
    code = """
def test_something():
    assert helper_used() == 42

def helper_used():
    return 42

def helper_unused():
    return 123

USED_CONSTANT = "used"
UNUSED_CONSTANT = "unused"

def test_another():
    assert USED_CONSTANT == "used"
"""

    expected = """
def test_something():
    assert helper_used() == 42

def helper_used():
    return 42

USED_CONSTANT = "used"

def test_another():
    assert USED_CONSTANT == "used"
"""

    result = remove_unused_definitions_from_pytest_file(cst.parse_module(code)).code
    # Normalize whitespace for comparison
    assert result.strip() == expected.strip()


def test_transitive_dependencies() -> None:
    """Test that dependencies of dependencies are preserved."""
    code = """
def test_main():
    assert direct_dependency() == 123

def direct_dependency():
    return indirect_dependency() + 100

def indirect_dependency():
    return 23

def completely_unused():
    return "this should be removed"
"""

    expected = """
def test_main():
    assert direct_dependency() == 123

def direct_dependency():
    return indirect_dependency() + 100

def indirect_dependency():
    return 23
"""

    result = remove_unused_definitions_from_pytest_file(cst.parse_module(code)).code
    # Normalize whitespace for comparison
    assert result.strip() == expected.strip()


def test_class_handling() -> None:
    """Test handling of classes and their dependencies."""
    code = """
class UnusedClass:
    def method(self):
        return 42

class UsedClass:
    def __init__(self):
        self.value = helper_function()

    def get_value(self):
        return self.value

def helper_function():
    return 100

def test_class_usage():
    obj = UsedClass()
    assert obj.get_value() == 100
"""

    expected = """
class UsedClass:
    def __init__(self):
        self.value = helper_function()

    def get_value(self):
        return self.value

def helper_function():
    return 100

def test_class_usage():
    obj = UsedClass()
    assert obj.get_value() == 100
"""

    result = remove_unused_definitions_from_pytest_file(cst.parse_module(code)).code
    # Normalize whitespace for comparison
    assert result.strip() == expected.strip()


def test_nested_function_dependencies() -> None:
    """Test that nested functions don't affect dependency tracking."""
    code = """
def test_outer_function():
    assert main_function() == 42

def main_function():
    # This inner function shouldn't be tracked as a separate entity
    def inner_helper():
        return 20

    return inner_helper() + external_helper()

def external_helper():
    return 22

def unused_function():
    return "This should be removed"
"""

    expected = """
def test_outer_function():
    assert main_function() == 42

def main_function():
    # This inner function shouldn't be tracked as a separate entity
    def inner_helper():
        return 20

    return inner_helper() + external_helper()

def external_helper():
    return 22
"""

    result = remove_unused_definitions_from_pytest_file(cst.parse_module(code)).code
    assert result.strip() == expected.strip()


def test_imported_but_unused() -> None:
    """Test that imports are not affected by the removal process."""
    code = """
import sys
import os
from pathlib import Path

def test_something():
    assert helper() == "test"

def helper():
    return "test"

def unused():
    return sys.version
"""

    expected = """
import sys
import os
from pathlib import Path

def test_something():
    assert helper() == "test"

def helper():
    return "test"
"""

    result = remove_unused_definitions_from_pytest_file(cst.parse_module(code)).code
    assert result.strip() == expected.strip()


def test_class_method_dependencies() -> None:
    """Test that class methods properly track external dependencies."""
    code = """
def test_class_usage():
    obj = UsedClass()
    assert obj.process_data() == "processed data"

class UsedClass:
    def process_data(self):
        return external_helper("data")

def external_helper(data):
    return f"processed {data}"

class UnusedClass:
    def __init__(self):
        self.helper = external_helper

    def do_something(self):
        return self.helper("unused")
"""

    expected = """
def test_class_usage():
    obj = UsedClass()
    assert obj.process_data() == "processed data"

class UsedClass:
    def process_data(self):
        return external_helper("data")

def external_helper(data):
    return f"processed {data}"
"""

    result = remove_unused_definitions_from_pytest_file(cst.parse_module(code)).code
    assert result.strip() == expected.strip()


def test_multiple_test_functions() -> None:
    """Test that all test functions are properly considered."""
    code = """
def test_first():
    assert helper1() == 1

def test_second():
    assert helper2() == 2

def helper1():
    return 1

def helper2():
    return 2

def helper3():
    # This is used indirectly by test_third
    return 3

def test_third():
    assert helper3() == 3

def unused_helper():
    return 4
"""

    expected = """
def test_first():
    assert helper1() == 1

def test_second():
    assert helper2() == 2

def helper1():
    return 1

def helper2():
    return 2

def helper3():
    # This is used indirectly by test_third
    return 3

def test_third():
    assert helper3() == 3
"""

    result = remove_unused_definitions_from_pytest_file(cst.parse_module(code)).code
    assert result.strip() == expected.strip()


def test_class_inheritance() -> None:
    """Test that class inheritance is properly handled."""
    code = """
def test_derived_class():
    derived = DerivedClass()
    assert derived.get_value() == "base-derived"

class BaseClass:
    def get_base_value(self):
        return "base"

class DerivedClass(BaseClass):
    def get_value(self):
        return self.get_base_value() + "-derived"

class UnusedClass(BaseClass):
    def get_unused_value(self):
        return "unused"
"""

    expected = """
def test_derived_class():
    derived = DerivedClass()
    assert derived.get_value() == "base-derived"

class BaseClass:
    def get_base_value(self):
        return "base"

class DerivedClass(BaseClass):
    def get_value(self):
        return self.get_base_value() + "-derived"
"""

    result = remove_unused_definitions_from_pytest_file(cst.parse_module(code)).code
    assert result.strip() == expected.strip()


def test_circular_dependencies() -> None:
    """Test that circular dependencies between definitions are handled properly."""
    code = """
def test_function():
    assert function_a() == "a calls b, b calls a"

def function_a():
    return "a calls b, " + function_b()

def function_b():
    return "b calls a"

def unused_circular_a():
    return "unused a calls " + unused_circular_b()

def unused_circular_b():
    return "unused b which calls " + unused_circular_a()
"""

    expected = """
def test_function():
    assert function_a() == "a calls b, b calls a"

def function_a():
    return "a calls b, " + function_b()

def function_b():
    return "b calls a"
"""

    result = remove_unused_definitions_from_pytest_file(cst.parse_module(code)).code
    assert result.strip() == expected.strip()


def test_dependencies_on_nonexistent_names() -> None:
    """Test behavior when code references names that aren't defined at top level."""
    code = """
def test_with_external_dep():
    assert uses_external() == True

def uses_external():
    # This references something not defined at top level in this file
    # Maybe it's imported or defined elsewhere
    return external_module.function()

def unused_function():
    return "This should be removed"
"""

    expected = """
def test_with_external_dep():
    assert uses_external() == True

def uses_external():
    # This references something not defined at top level in this file
    # Maybe it's imported or defined elsewhere
    return external_module.function()
"""

    result = remove_unused_definitions_from_pytest_file(cst.parse_module(code)).code
    assert result.strip() == expected.strip()


def test_complex_assignments() -> None:
    """Test handling of complex variable assignments."""
    code = """
def test_variable_usage():
    assert USED_TUPLE[0] == 1

# Different assignment patterns
USED_TUPLE = (1, 2, 3)
USED_DICT = {"key": "value"}
USED_VAL = 0
USED_VAL += 5
UNUSED_LIST = [4, 5, 6]
A, B = 1, 2

def test_another():
    USED_VAL += 3
    assert USED_DICT["key"] == "value"
"""

    expected = """
def test_variable_usage():
    assert USED_TUPLE[0] == 1

# Different assignment patterns
USED_TUPLE = (1, 2, 3)
USED_DICT = {"key": "value"}
USED_VAL = 0
USED_VAL += 5

def test_another():
    USED_VAL += 3
    assert USED_DICT["key"] == "value"
"""

    result = remove_unused_definitions_from_pytest_file(cst.parse_module(code)).code
    assert result.strip() == expected.strip()


def test_only_test_functions() -> None:
    """Test that nothing is removed when only test functions are present."""
    code = """
def test_one():
    assert 1 == 1

def test_two():
    assert 2 == 2
"""
    expected = code
    result = remove_unused_definitions_from_pytest_file(cst.parse_module(code)).code
    assert result.strip() == expected.strip()


def test_recursive_function() -> None:
    """Test that recursive functions are properly handled."""
    code = """
def test_factorial():
    assert factorial(5) == 120

def factorial(n):
    # This function calls itself recursively
    if n <= 1:
        return 1
    return n * factorial(n - 1)

def unused_recursive_function(n):
    # This recursive function is not used by any test
    if n <= 0:
        return 0
    return n + unused_recursive_function(n - 1)
"""

    expected = """
def test_factorial():
    assert factorial(5) == 120

def factorial(n):
    # This function calls itself recursively
    if n <= 1:
        return 1
    return n * factorial(n - 1)
"""

    result = remove_unused_definitions_from_pytest_file(cst.parse_module(code)).code
    assert result.strip() == expected.strip()


def test_tuple_unpacking_with_used_variables() -> None:
    """Test that tuple unpacking preserves used variables."""
    code = """
def test_function():
    assert USED_VAR == 1

# Tuple unpacking with one used variable
USED_VAR, UNUSED_VAR = 1, 2

# Nested unpacking with a used variable
A, (B, USED_NESTED) = 1, (2, 3)

def test_nested():
    assert USED_NESTED == 3
"""

    expected = """
def test_function():
    assert USED_VAR == 1

# Tuple unpacking with one used variable
USED_VAR, UNUSED_VAR = 1, 2

# Nested unpacking with a used variable
A, (B, USED_NESTED) = 1, (2, 3)

def test_nested():
    assert USED_NESTED == 3
"""

    result = remove_unused_definitions_from_pytest_file(cst.parse_module(code)).code
    assert result.strip() == expected.strip()


def test_tuple_unpacking_with_used_variables() -> None:
    """Test that tuple unpacking preserves used variables."""
    code = """
def test_function():
    assert USED_VAR == 1

# Tuple unpacking with one used variable
USED_VAR, UNUSED_VAR = 1, 2

# Nested unpacking with a used variable
A, (B, USED_NESTED) = 1, (2, 3)

def test_nested():
    assert USED_NESTED == 3
"""

    expected = """
def test_function():
    assert USED_VAR == 1

# Tuple unpacking with one used variable
USED_VAR, UNUSED_VAR = 1, 2

# Nested unpacking with a used variable
A, (B, USED_NESTED) = 1, (2, 3)

def test_nested():
    assert USED_NESTED == 3
"""

    result = remove_unused_definitions_from_pytest_file(cst.parse_module(code)).code
    assert result.strip() == expected.strip()


def test_annotated_assignments() -> None:
    """Test handling of type-annotated assignments."""
    code = """
def test_annotated():
    assert USED_ANNOTATED == 42

# Annotated assignments
USED_ANNOTATED: int = 42
UNUSED_ANNOTATED: str = "hello"

# Annotated variables individually
USED_X: int = 1
UNUSED_Y: int = 2

def test_another():
    assert USED_X == 1
    assert USED_ANNOTATED == 42
"""

    expected = """
def test_annotated():
    assert USED_ANNOTATED == 42

# Annotated assignments
USED_ANNOTATED: int = 42

# Annotated variables individually
USED_X: int = 1

def test_another():
    assert USED_X == 1
    assert USED_ANNOTATED == 42
"""

    result = remove_unused_definitions_from_pytest_file(cst.parse_module(code)).code
    assert result.strip() == expected.strip()


def test_augmented_assignments() -> None:
    """Test handling of augmented assignments."""
    code = """
def test_counter():
    assert COUNTER == 10

# Regular assignment
COUNTER = 0
UNUSED_COUNTER = 0

def increment_counter():
    global COUNTER
    # Augmented assignment
    COUNTER += 10

def test_setup():
    increment_counter()
"""

    expected = """
def test_counter():
    assert COUNTER == 10

# Regular assignment
COUNTER = 0

def increment_counter():
    global COUNTER
    # Augmented assignment
    COUNTER += 10

def test_setup():
    increment_counter()
"""

    result = remove_unused_definitions_from_pytest_file(cst.parse_module(code)).code
    assert result.strip() == expected.strip()


def test_multiple_assignment() -> None:
    """Test handling of multiple variable assignment in a single statement."""
    code = """
def test_function():
    assert X == 5

# Multiple assignment
X = Y = Z = 5

def test_unused():
    # Note: Y and Z are never referenced, only X is used
    pass
"""

    expected = """
def test_function():
    assert X == 5

# Multiple assignment
X = Y = Z = 5

def test_unused():
    # Note: Y and Z are never referenced, only X is used
    pass
"""

    result = remove_unused_definitions_from_pytest_file(cst.parse_module(code)).code
    assert result.strip() == expected.strip()


def test_list_and_dict_unpacking() -> None:
    """Test unpacking with lists and more complex structures."""
    code = """
def test_unpacking():
    assert FIRST == 1

# List unpacking
[FIRST, SECOND] = [1, 2]
[UNUSED_A, UNUSED_B] = [3, 4]

# Dict items - this is not tracked as a dependency of the test
# because it's used inside a loop structure
ITEMS = {"a": 1, "b": 2}.items()
for KEY, VALUE in ITEMS:
    # Using unpacked variables in a loop
    pass

# This variable will be preserved because it's directly referenced by a test
LOOP_TEST = 42

def helper():
    # This function uses SECOND
    return SECOND

def test_helper():
    assert helper() == 2

def test_loop_var():
    assert LOOP_TEST == 42
    # Note: We cannot reference ITEMS here, as dependencies inside
    # loops or control flow structures are not tracked
"""

    expected = """
def test_unpacking():
    assert FIRST == 1

# List unpacking
[FIRST, SECOND] = [1, 2]

# This variable will be preserved because it's directly referenced by a test
LOOP_TEST = 42

def helper():
    # This function uses SECOND
    return SECOND

def test_helper():
    assert helper() == 2

def test_loop_var():
    assert LOOP_TEST == 42
    # Note: We cannot reference ITEMS here, as dependencies inside
    # loops or control flow structures are not tracked
"""

    result = remove_unused_definitions_from_pytest_file(cst.parse_module(code)).code
    assert result.strip() == expected.strip()


def test_augmented_assignments() -> None:
    """Test handling of augmented assignments."""
    code = """
def test_counter():
    assert COUNTER == 10

# Regular assignment
COUNTER = 0
UNUSED_COUNTER = 0

def increment_counter():
    global COUNTER
    # Augmented assignment
    COUNTER += 10

def test_setup():
    increment_counter()
"""

    expected = """
def test_counter():
    assert COUNTER == 10

# Regular assignment
COUNTER = 0

def increment_counter():
    global COUNTER
    # Augmented assignment
    COUNTER += 10

def test_setup():
    increment_counter()
"""

    result = remove_unused_definitions_from_pytest_file(cst.parse_module(code)).code
    assert result.strip() == expected.strip()


def test_multiple_assignment() -> None:
    """Test handling of multiple variable assignment in a single statement."""
    code = """
def test_function():
    assert X == 5

# Multiple assignment
X = Y = Z = 5

def test_unused():
    # Note: Y and Z are never referenced, only X is used
    pass
"""

    expected = """
def test_function():
    assert X == 5

# Multiple assignment
X = Y = Z = 5

def test_unused():
    # Note: Y and Z are never referenced, only X is used
    pass
"""

    result = remove_unused_definitions_from_pytest_file(cst.parse_module(code)).code
    assert result.strip() == expected.strip()


def test_if_else_statements() -> None:
    """Test handling of if-else statements and their dependencies."""
    code = """
def test_conditional():
    assert conditional_function(True) == "condition true"
    assert conditional_function(False) == "condition false"

def conditional_function(condition):
    if condition:
        return true_path()
    else:
        return false_path()

def true_path():
    return "condition true"

def false_path():
    return "condition false"

def unused_conditional(flag):
    if flag:
        return "unused true"
    else:
        return "unused false"
"""

    expected = """
def test_conditional():
    assert conditional_function(True) == "condition true"
    assert conditional_function(False) == "condition false"

def conditional_function(condition):
    if condition:
        return true_path()
    else:
        return false_path()

def true_path():
    return "condition true"

def false_path():
    return "condition false"
"""

    result = remove_unused_definitions_from_pytest_file(cst.parse_module(code)).code
    assert result.strip() == expected.strip()


def test_try_except_finally() -> None:
    """Test handling of try-except-finally blocks and their dependencies."""
    code = """
def test_exception_handling():
    assert exception_handler(10) == "handled: 10"
    assert exception_handler("test") == "caught TypeError"

def exception_handler(value):
    try:
        return "handled: " + str(process_value(value))
    except TypeError:
        return "caught TypeError"
    finally:
        cleanup()

def process_value(value):
    return value + 5

def cleanup():
    # Cleanup logic
    pass

def unused_exception_function():
    try:
        return "unused try"
    except:
        return "unused except"
    finally:
        pass
"""

    expected = """
def test_exception_handling():
    assert exception_handler(10) == "handled: 10"
    assert exception_handler("test") == "caught TypeError"

def exception_handler(value):
    try:
        return "handled: " + str(process_value(value))
    except TypeError:
        return "caught TypeError"
    finally:
        cleanup()

def process_value(value):
    return value + 5

def cleanup():
    # Cleanup logic
    pass
"""

    result = remove_unused_definitions_from_pytest_file(cst.parse_module(code)).code
    assert result.strip() == expected.strip()


def test_complex_conditional_dependency() -> None:
    """Test complex conditional dependencies where some code paths may be unused."""
    code = """
def test_branches():
    # Only tests path A and C, not B
    assert complex_branches("A") == "path A"
    assert complex_branches("C") == "path C"

def complex_branches(path):
    if path == "A":
        return path_a()
    elif path == "B":
        return path_b()
    else:
        return path_c()

def path_a():
    return "path A"

def path_b():
    return "path B"

def path_c():
    return "path C"

def unused_path():
    return "unused path"
"""

    expected = """
def test_branches():
    # Only tests path A and C, not B
    assert complex_branches("A") == "path A"
    assert complex_branches("C") == "path C"

def complex_branches(path):
    if path == "A":
        return path_a()
    elif path == "B":
        return path_b()
    else:
        return path_c()

def path_a():
    return "path A"

def path_b():
    return "path B"

def path_c():
    return "path C"
"""

    result = remove_unused_definitions_from_pytest_file(cst.parse_module(code)).code
    assert result.strip() == expected.strip()


def test_top_level_if_else() -> None:
    """Test handling of top-level if-else with variable and function definitions."""
    code = """
import sys

# Top-level if-else controlling what gets defined
if sys.platform == 'win32':
    # Windows-specific code
    PLATFORM = "windows"

    def get_path_separator():
        return '\\\\'
else:
    # Non-windows code
    PLATFORM = "other"

    def get_path_separator():
        return '/'

# Unused variable defined outside if-else
UNUSED_VAR = "not used"

def test_platform_info():
    assert PLATFORM in ["windows", "other"]
    assert get_path_separator() in ['\\\\', '/']
"""

    expected = """
import sys

# Top-level if-else controlling what gets defined
if sys.platform == 'win32':
    # Windows-specific code
    PLATFORM = "windows"

    def get_path_separator():
        return '\\\\'
else:
    # Non-windows code
    PLATFORM = "other"

    def get_path_separator():
        return '/'

def test_platform_info():
    assert PLATFORM in ["windows", "other"]
    assert get_path_separator() in ['\\\\', '/']
"""

    result = remove_unused_definitions_from_pytest_file(cst.parse_module(code)).code
    assert result.strip() == expected.strip()


def test_top_level_try_except() -> None:
    """Test handling of top-level try-except with variable and function definitions."""
    code = """
# Top-level try-except that defines variables and functions
try:
    import numpy as np
    NUMPY_AVAILABLE = True

    def calculate_mean(values):
        return np.mean(values)
except ImportError:
    NUMPY_AVAILABLE = False

    def calculate_mean(values):
        return sum(values) / len(values)

# Unused function defined outside try-except
def unused_function():
    return "This is not used"

def test_mean_calculation():
    assert calculate_mean([1, 2, 3, 4]) == 2.5
    assert NUMPY_AVAILABLE in [True, False]
"""

    expected = """
# Top-level try-except that defines variables and functions
try:
    import numpy as np
    NUMPY_AVAILABLE = True

    def calculate_mean(values):
        return np.mean(values)
except ImportError:
    NUMPY_AVAILABLE = False

    def calculate_mean(values):
        return sum(values) / len(values)

def test_mean_calculation():
    assert calculate_mean([1, 2, 3, 4]) == 2.5
    assert NUMPY_AVAILABLE in [True, False]
"""

    result = remove_unused_definitions_from_pytest_file(cst.parse_module(code)).code
    assert result.strip() == expected.strip()


def test_mixed_top_level_control_flow() -> None:
    """Test handling of mixed top-level control structures with different definitions."""
    code = """
import os

# Top-level if-else
if os.name == 'posix':
    CURRENT_OS = 'posix'
else:
    CURRENT_OS = 'other'

# Top-level try-except
try:
    from math import tau
    TAU_AVAILABLE = True
except ImportError:
    TAU_AVAILABLE = False
    tau = 6.28318

# Unused top-level variable
UNUSED_CONSTANT = "Not used anywhere"

# Used function that references variables from control structures
def get_circle_circumference(radius):
    return tau * radius

# Unused function
def unused_helper():
    return "Not called by tests"

def test_os_detection():
    assert CURRENT_OS in ['posix', 'other']

def test_circumference():
    assert get_circle_circumference(1) > 6.28
    assert get_circle_circumference(1) < 6.29
"""

    expected = """
import os

# Top-level if-else
if os.name == 'posix':
    CURRENT_OS = 'posix'
else:
    CURRENT_OS = 'other'

# Top-level try-except
try:
    from math import tau
except ImportError:
    tau = 6.28318

# Used function that references variables from control structures
def get_circle_circumference(radius):
    return tau * radius

def test_os_detection():
    assert CURRENT_OS in ['posix', 'other']

def test_circumference():
    assert get_circle_circumference(1) > 6.28
    assert get_circle_circumference(1) < 6.29
"""

    result = remove_unused_definitions_from_pytest_file(cst.parse_module(code)).code
    assert result.strip() == expected.strip()


def test_top_level_with_context_manager() -> None:
    """Test handling of top-level 'with' statements and their variable definitions."""
    code = """
import tempfile
import os

# Top-level with statement defining variables
with tempfile.TemporaryDirectory() as temp_dir:
    TEST_DIR = temp_dir
    TEST_FILE = os.path.join(temp_dir, "test.txt")
    UNUSED_PATH = os.path.join(temp_dir, "unused.txt")

# Create a file handler that isn't used
with open(os.devnull, "w") as f:
    NULL_HANDLER = f
    UNUSED_HANDLER = f

def get_test_file_path():
    return TEST_FILE

def unused_function():
    return UNUSED_PATH

def test_file_path():
    assert TEST_DIR is not None
    assert get_test_file_path().endswith("test.txt")
"""

    expected = """
import tempfile
import os

# Top-level with statement defining variables
with tempfile.TemporaryDirectory() as temp_dir:
    TEST_DIR = temp_dir
    TEST_FILE = os.path.join(temp_dir, "test.txt")

def get_test_file_path():
    return TEST_FILE

def test_file_path():
    assert TEST_DIR is not None
    assert get_test_file_path().endswith("test.txt")
"""

    result = remove_unused_definitions_from_pytest_file(cst.parse_module(code)).code
    assert result.strip() == expected.strip()


def test_ultralytics() -> None:
    code = """
import glob
import math
import os
from pathlib import Path

import cv2
import numpy as np
# imports
import pytest  # used for our unit tests
from PIL import Image
from ultralytics.data.loaders import LoadImagesAndVideos
from ultralytics.data.utils import FORMATS_HELP_MSG, IMG_FORMATS, VID_FORMATS
from ultralytics.utils import LOGGER
from ultralytics.utils.checks import check_requirements
from ultralytics.utils.patches import imread

# Define constants
LOADERS = (LoadStreams, LoadPilAndNumpy, LoadImagesAndVideos, LoadScreenshots)


# unit tests

def test_single_image_file():
    loader = LoadImagesAndVideos(["image1.jpg"], batch=1)

def test_single_image_file_large_batch():
    loader = LoadImagesAndVideos(["image1.jpg"], batch=5)
"""
    expected = """
import glob
import math
import os
from pathlib import Path

import cv2
import numpy as np
# imports
import pytest  # used for our unit tests
from PIL import Image
from ultralytics.data.loaders import LoadImagesAndVideos
from ultralytics.data.utils import FORMATS_HELP_MSG, IMG_FORMATS, VID_FORMATS
from ultralytics.utils import LOGGER
from ultralytics.utils.checks import check_requirements
from ultralytics.utils.patches import imread


# unit tests

def test_single_image_file():
    loader = LoadImagesAndVideos(["image1.jpg"], batch=1)

def test_single_image_file_large_batch():
    loader = LoadImagesAndVideos(["image1.jpg"], batch=5)
"""
    result = remove_unused_definitions_from_pytest_file(cst.parse_module(code)).code
    assert result.strip() == expected.strip()


def test_unittest_testcase_subclass() -> None:
    """Test that classes inheriting from unittest.TestCase and their dependencies are preserved."""
    code = """
import unittest

def sorter(arr):
    # This function should be preserved because it's used by TestSorter
    return sorted(arr)

def unused_function():
    # This should be removed
    return "unused"

class TestSorter(unittest.TestCase):
    def test_empty_list(self):
        self.assertEqual(sorter([]), [])

    def test_single_element(self):
        self.assertEqual(sorter([42]), [42])

class UnusedClass:
    def method(self):
        return "unused"
"""

    expected = """
import unittest

def sorter(arr):
    # This function should be preserved because it's used by TestSorter
    return sorted(arr)

class TestSorter(unittest.TestCase):
    def test_empty_list(self):
        self.assertEqual(sorter([]), [])

    def test_single_element(self):
        self.assertEqual(sorter([42]), [42])
"""

    result = remove_unused_definitions_from_pytest_file(cst.parse_module(code)).code
    assert result.strip() == expected.strip()


def test_abstract_base_class_inheritance() -> None:
    """Test that abstract base classes used only for inheritance are preserved.

    This mimics the real-world case where LLM generates mock classes that inherit
    from a base class, and the base class should not be removed even though it's
    only referenced in the inheritance declaration.
    """
    code = """
class LayoutDumper:
    layout_source: str = "unknown"
    def dump(self) -> dict:
        raise NotImplementedError()

class ObjectDetectionLayoutDumper(LayoutDumper):
    def __init__(self, layout):
        self._layout = layout
    def dump(self) -> dict:
        return self._layout

class ExtractedLayoutDumper(LayoutDumper):
    def __init__(self, layout):
        self._layout = layout
    def dump(self) -> dict:
        return self._layout

class UnusedClass:
    pass

def test_object_detection():
    dumper = ObjectDetectionLayoutDumper({})
    assert dumper.dump() == {}

def test_extracted():
    dumper = ExtractedLayoutDumper({"text": "hello"})
    assert dumper.dump() == {"text": "hello"}
"""

    expected = """
class LayoutDumper:
    layout_source: str = "unknown"
    def dump(self) -> dict:
        raise NotImplementedError()

class ObjectDetectionLayoutDumper(LayoutDumper):
    def __init__(self, layout):
        self._layout = layout
    def dump(self) -> dict:
        return self._layout

class ExtractedLayoutDumper(LayoutDumper):
    def __init__(self, layout):
        self._layout = layout
    def dump(self) -> dict:
        return self._layout

def test_object_detection():
    dumper = ObjectDetectionLayoutDumper({})
    assert dumper.dump() == {}

def test_extracted():
    dumper = ExtractedLayoutDumper({"text": "hello"})
    assert dumper.dump() == {"text": "hello"}
"""

    result = remove_unused_definitions_from_pytest_file(cst.parse_module(code)).code
    assert result.strip() == expected.strip()
