import ast
import os


def check_for_litellm_module_deletion(base_dir):
    """
    Checks for code patterns that delete litellm modules from sys.modules
    in the test_litellm directory.

    Specifically looks for patterns like:
    for module in list(sys.modules.keys()):
        if module.startswith("litellm"):
            del sys.modules[module]
    """
    problematic_files = []
    test_dir = os.path.join(base_dir, "test_litellm")

    if not os.path.exists(test_dir):
        print(f"Warning: Directory {test_dir} does not exist.")
        return []

    print(f"Checking directory: {test_dir}")

    for root, _, files in os.walk(test_dir):
        for file in files:
            if file.endswith(".py"):
                file_path = os.path.join(root, file)
                try:
                    with open(file_path, "r") as f:
                        tree = ast.parse(f.read())
                except SyntaxError:
                    print(f"Warning: Syntax error in file {file_path}")
                    continue

                # Check for litellm module deletion patterns
                if has_litellm_module_deletion(tree):
                    relative_path = os.path.relpath(file_path, base_dir)
                    problematic_files.append(relative_path)
                    print(f"Found litellm module deletion in: {relative_path}")

    return problematic_files


def has_litellm_module_deletion(tree):
    """
    Checks if the AST contains patterns that delete litellm modules from sys.modules.

    Looks for:
    1. Loops over sys.modules.keys()
    2. Conditions checking if module startswith "litellm"
    3. del sys.modules[module] statements
    """

    class LiteLLMDeletionVisitor(ast.NodeVisitor):
        def __init__(self):
            self.has_sys_modules_loop = False
            self.has_litellm_check = False
            self.has_del_sys_modules = False
            self.current_for_target = None

        def visit_For(self, node):
            # Check if we're looping over sys.modules.keys()
            if (
                isinstance(node.iter, ast.Call)
                and isinstance(node.iter.func, ast.Attribute)
                and isinstance(node.iter.func.value, ast.Attribute)
                and isinstance(node.iter.func.value.value, ast.Name)
                and node.iter.func.value.value.id == "sys"
                and node.iter.func.value.attr == "modules"
                and node.iter.func.attr == "keys"
            ):

                self.has_sys_modules_loop = True
                if isinstance(node.target, ast.Name):
                    self.current_for_target = node.target.id

                # Check the body of the for loop
                for stmt in node.body:
                    self.visit(stmt)

            # Also check for list(sys.modules.keys()) pattern
            elif (
                isinstance(node.iter, ast.Call)
                and isinstance(node.iter.func, ast.Name)
                and node.iter.func.id == "list"
                and len(node.iter.args) == 1
                and isinstance(node.iter.args[0], ast.Call)
                and isinstance(node.iter.args[0].func, ast.Attribute)
                and isinstance(node.iter.args[0].func.value, ast.Attribute)
                and isinstance(node.iter.args[0].func.value.value, ast.Name)
                and node.iter.args[0].func.value.value.id == "sys"
                and node.iter.args[0].func.value.attr == "modules"
                and node.iter.args[0].func.attr == "keys"
            ):

                self.has_sys_modules_loop = True
                if isinstance(node.target, ast.Name):
                    self.current_for_target = node.target.id

                # Check the body of the for loop
                for stmt in node.body:
                    self.visit(stmt)

            self.generic_visit(node)

        def visit_If(self, node):
            # Check for conditions like module.startswith("litellm")
            if (
                isinstance(node.test, ast.Call)
                and isinstance(node.test.func, ast.Attribute)
                and isinstance(node.test.func.value, ast.Name)
                and node.test.func.value.id == self.current_for_target
                and node.test.func.attr == "startswith"
                and len(node.test.args) == 1
                and isinstance(node.test.args[0], ast.Constant)
                and node.test.args[0].value == "litellm"
            ):

                self.has_litellm_check = True

                # Check the body of the if statement
                for stmt in node.body:
                    self.visit(stmt)

            self.generic_visit(node)

        def visit_Delete(self, node):
            # Check for del sys.modules[module]
            for target in node.targets:
                if (
                    isinstance(target, ast.Subscript)
                    and isinstance(target.value, ast.Attribute)
                    and isinstance(target.value.value, ast.Name)
                    and target.value.value.id == "sys"
                    and target.value.attr == "modules"
                    and isinstance(target.slice, ast.Name)
                    and target.slice.id == self.current_for_target
                ):

                    self.has_del_sys_modules = True

            self.generic_visit(node)

    visitor = LiteLLMDeletionVisitor()
    visitor.visit(tree)

    return (
        visitor.has_sys_modules_loop
        and visitor.has_litellm_check
        and visitor.has_del_sys_modules
    )


def main():
    """
    Main function to check for litellm module deletion patterns in test files.
    """
    # local dir
    # tests_dir = "../../tests/"

    # ci/cd dir
    tests_dir = "./tests/"

    problematic_files = check_for_litellm_module_deletion(tests_dir)

    if problematic_files:
        print("\nERROR: Found files that delete litellm modules from sys.modules:")
        for file_path in problematic_files:
            print(f"  - {file_path}")

        raise Exception(
            f"Found {len(problematic_files)} file(s) that delete litellm modules from sys.modules. "
            f"This can cause import issues and test failures. Files: {problematic_files}"
        )
    else:
        print("✓ No litellm module deletion patterns found in test_litellm directory.")


if __name__ == "__main__":
    main()