import ast
import os

ALLOWED_FILES = [
    # local files
    "../../litellm/litellm_core_utils/litellm.logging_callback_manager.py",
    "../../litellm/proxy/common_utils/callback_utils.py",
    # when running on ci/cd
    "./litellm/litellm_core_utils/litellm.logging_callback_manager.py",
    "./litellm/proxy/common_utils/callback_utils.py",
]

warning_msg = "this is a serious violation. Callbacks must only be modified through LoggingCallbackManager"


def check_for_callback_modifications(file_path):
    """
    Checks if any direct modifications to specific litellm callback lists are made in the given file.
    Also prints the violating line of code.
    """
    print("..checking file=", file_path)
    if file_path in ALLOWED_FILES:
        return []

    violations = []
    with open(file_path, "r") as file:
        try:
            lines = file.readlines()
            tree = ast.parse("".join(lines))
        except SyntaxError:
            print(f"Warning: Syntax error in file {file_path}")
            return violations

    protected_lists = [
        "callbacks",
        "success_callback",
        "failure_callback",
        "_async_success_callback",
        "_async_failure_callback",
    ]

    forbidden_operations = ["append", "extend", "insert"]

    for node in ast.walk(tree):
        # Check for attribute calls like litellm.callbacks.append()
        if isinstance(node, ast.Call) and isinstance(node.func, ast.Attribute):
            # Get the full attribute chain
            attr_chain = []
            current = node.func
            while isinstance(current, ast.Attribute):
                attr_chain.append(current.attr)
                current = current.value
            if isinstance(current, ast.Name):
                attr_chain.append(current.id)

            # Reverse to get the chain from root to leaf
            attr_chain = attr_chain[::-1]

            # Check if the attribute chain starts with 'litellm' and modifies a protected list
            if (
                len(attr_chain) >= 3
                and attr_chain[0] == "litellm"
                and attr_chain[2] in forbidden_operations
            ):
                protected_list = attr_chain[1]
                operation = attr_chain[2]
                if (
                    protected_list in protected_lists
                    and operation in forbidden_operations
                ):
                    violating_line = lines[node.lineno - 1].strip()
                    violations.append(
                        f"Found violation in file {file_path} line {node.lineno}: '{violating_line}'. "
                        f"Direct modification of 'litellm.{protected_list}' using '{operation}' is not allowed. "
                        f"Please use LoggingCallbackManager instead. {warning_msg}"
                    )

    return violations


def scan_directory_for_callback_modifications(base_dir):
    """
    Scans all Python files in the directory tree for unauthorized callback list modifications.
    """
    all_violations = []
    for root, _, files in os.walk(base_dir):
        for file in files:
            if file.endswith(".py"):
                file_path = os.path.join(root, file)
                violations = check_for_callback_modifications(file_path)
                all_violations.extend(violations)
    return all_violations


def test_no_unauthorized_callback_modifications():
    """
    Test to ensure callback lists are not modified directly anywhere in the codebase.
    """
    base_dir = "./litellm"  # Adjust this path as needed
    # base_dir = "../../litellm"  # LOCAL TESTING

    violations = scan_directory_for_callback_modifications(base_dir)
    if violations:
        print(f"\nFound {len(violations)} callback modification violations:")
        for violation in violations:
            print("\n" + violation)
        raise AssertionError(
            "Found unauthorized callback modifications. See above for details."
        )


if __name__ == "__main__":
    test_no_unauthorized_callback_modifications()