"""
Code quality check: Ensure _get_model_cost_key only uses O(1) operations.

Simple pattern-based check for O(n) operations in _get_model_cost_key.
"""

import re
import os


def _function_has_on_operations(all_lines, func_name, visited=None):
    """
    Check if a function contains O(n) operations by searching for it in the file.
    Recursively checks called functions as well.
    """
    if visited is None:
        visited = set()

    # Prevent infinite recursion
    if func_name in visited:
        return False
    visited.add(func_name)

    func_start = None
    func_end = None

    for i, line in enumerate(all_lines):
        if func_start is None and f"def {func_name}(" in line:
            func_start = i
        elif func_start is not None:
            # Function ends when we hit next def at module level
            if (
                line.strip()
                and not line.startswith(" ")
                and not line.startswith("\t")
                and line.startswith("def ")
            ):
                func_end = i
                break

    if func_start is None or func_end is None:
        return False

    # Check function body for O(n) patterns
    func_lines = all_lines[func_start:func_end]

    for line in func_lines:
        # Skip comments and docstrings
        line_stripped = line.strip()
        if (
            line_stripped.startswith("#")
            or line_stripped.startswith('"""')
            or line_stripped.startswith("'''")
        ):
            continue

        # Check for for loops
        if re.search(r"\bfor\s+\w+\s+in\s+", line):
            return True
        # Check for while loops
        if re.search(r"\bwhile\s+", line):
            return True
        # Check for comprehensions
        if re.search(r"\[.*\s+for\s+.*\s+in\s+", line) or re.search(
            r"\{.*\s+for\s+.*\s+in\s+", line
        ):
            return True

        # Recursively check called functions (check all, don't skip any in recursive checks)
        func_call_match = re.search(r"\b([a-z_][a-z0-9_]*)\s*\(", line)
        if func_call_match:
            called_func = func_call_match.group(1)
            if called_func.startswith("_"):
                if _function_has_on_operations(all_lines, called_func, visited):
                    return True

    return False


def check_get_model_cost_key_performance():
    """
    Check that _get_model_cost_key doesn't contain O(n) operations.
    """
    utils_file = "./litellm/utils.py"

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

    with open(utils_file, "r", encoding="utf-8") as f:
        lines = f.readlines()

    # Find the _get_model_cost_key function
    func_start = None
    func_end = None

    for i, line in enumerate(lines):
        if func_start is None and "def _get_model_cost_key(" in line:
            func_start = i
        elif func_start is not None:
            # Function ends when we hit next def at module level (no indentation)
            if (
                line.strip()
                and not line.startswith(" ")
                and not line.startswith("\t")
                and line.startswith("def ")
            ):
                func_end = i
                break

    if func_start is None:
        print("Warning: Could not find _get_model_cost_key function")
        return []

    if func_end is None:
        func_end = len(lines)

    # Extract function body
    func_lines = lines[func_start:func_end]
    problematic_lines = []

    # Track if we're inside a docstring
    in_docstring = False
    docstring_quote = None

    # Check for O(n) patterns
    for i, line in enumerate(func_lines, start=func_start + 1):
        line_stripped = line.strip()

        # Track docstring state (handle both single-line and multi-line docstrings)
        if not in_docstring:
            if line_stripped.startswith('"""') or line_stripped.startswith("'''"):
                docstring_quote = '"""' if line_stripped.startswith('"""') else "'''"
                # Check if it's a single-line docstring
                if line_stripped.count(docstring_quote) >= 2:
                    in_docstring = False  # Single-line, skip this line
                    continue
                else:
                    in_docstring = True
                    continue
        else:
            # Inside multi-line docstring, check for closing quote
            if docstring_quote is not None and docstring_quote in line:
                in_docstring = False
                docstring_quote = None
            continue  # Skip all lines inside docstring

        # Skip comments
        if line_stripped.startswith("#"):
            continue

        # Check for for loops
        if re.search(r"\bfor\s+\w+\s+in\s+", line):
            # Allow helper function calls (they're conditional)
            if not re.search(
                r"(_rebuild_model_cost_lowercase_map|_handle_stale_map_entry_rebuild|_handle_new_key_with_scan)",
                line,
            ):
                problematic_lines.append((i, "for loop", line_stripped))

        # Check for while loops
        if re.search(r"\bwhile\s+", line):
            problematic_lines.append((i, "while loop", line_stripped))

        # Check for comprehensions
        if re.search(r"\[.*\s+for\s+.*\s+in\s+", line) or re.search(
            r"\{.*\s+for\s+.*\s+in\s+", line
        ):
            problematic_lines.append((i, "comprehension", line_stripped))

        # Check for problematic function calls
        problematic_funcs = [
            "enumerate",
            "zip",
            "map",
            "filter",
            "sorted",
            "any",
            "all",
            "sum",
            "max",
            "min",
        ]
        for func in problematic_funcs:
            if re.search(rf"\b{func}\s*\(", line):
                problematic_lines.append((i, f"call to {func}()", line_stripped))

        # Check for calls to functions that might have O(n) operations
        # Allow known helper functions that are conditional
        allowed_helpers = [
            "_rebuild_model_cost_lowercase_map",
            "_handle_stale_map_entry_rebuild",
            "_handle_new_key_with_scan",
        ]

        # Check for function calls (pattern: function_name(...), but not function definitions)
        # Skip function definitions (def function_name(...))
        if not re.search(r"\bdef\s+", line):
            func_call_match = re.search(r"\b([a-z_][a-z0-9_]*)\s*\(", line)
            if func_call_match:
                func_name = func_call_match.group(1)
                # If it's a call to a function that might have O(n) operations, check it
                if func_name not in allowed_helpers and func_name.startswith("_"):
                    # Check if this function has O(n) operations
                    if _function_has_on_operations(lines, func_name):
                        problematic_lines.append(
                            (
                                i,
                                f"call to {func_name}() which contains O(n) operations",
                                line_stripped,
                            )
                        )

    return problematic_lines


def main():
    """Main function to check _get_model_cost_key performance requirements."""
    problematic_lines = check_get_model_cost_key_performance()

    if problematic_lines:
        print("\nERROR: Found O(n) operations in _get_model_cost_key:")
        for line_num, operation, context in problematic_lines:
            print(f"  Line {line_num}: {operation} - {context}")

        print(
            "\nWARNING: Only O(1) lookup operations are acceptable in _get_model_cost_key."
        )
        print("Any O(n) operations will cause severe CPU overhead.")

        raise Exception(
            f"Found {len(problematic_lines)} O(n) operation(s) in _get_model_cost_key. "
            f"This violates the performance requirement."
        )
    else:
        print(
            "OK: No O(n) operations found in _get_model_cost_key. Performance requirement satisfied."
        )


if __name__ == "__main__":
    main()
