import ast
import os
import re
from typing import List, Dict, Any


class SensitiveLogDetector(ast.NodeVisitor):
    """
    Detects logger.info() statements that might log sensitive request/response data.
    """

    def __init__(self):
        self.violations = []
        self.current_file = None

    def set_file(self, file_path: str):
        """Set the current file being analyzed"""
        self.current_file = file_path

    def visit_Call(self, node):
        """Visit function calls to detect logger.info() with sensitive data"""
        if self._is_logger_info_call(node):
            # Check all arguments to the logger.info() call
            for arg in node.args:
                if self._contains_sensitive_data(arg):
                    violation = {
                        "file": self.current_file,
                        "line": node.lineno,
                        "call": self._get_call_string(node),
                        "reason": self._get_violation_reason(arg),
                        "arg": self._get_arg_string(arg),
                    }
                    self.violations.append(violation)

        self.generic_visit(node)

    def _is_logger_info_call(self, node) -> bool:
        """Check if this is a logger.info() call"""
        if not isinstance(node.func, ast.Attribute):
            return False

        # Check for various logger patterns:
        # logger.info(), verbose_logger.info(), verbose_proxy_logger.info(), etc.
        if node.func.attr == "info":
            if isinstance(node.func.value, ast.Name):
                logger_name = node.func.value.id
                return any(
                    pattern in logger_name.lower() for pattern in ["logger", "log"]
                )

        return False

    def _contains_sensitive_data(self, arg) -> bool:
        """Check if the argument might contain sensitive data"""
        # Convert argument to string for analysis
        arg_str = self._get_arg_string(arg).lower()

        # Skip obvious non-sensitive patterns
        non_sensitive_patterns = [
            r'^["\'][\w\s\-_:.,!?]*["\']$',  # Simple static strings
            r'^["\'][^{%]*["\']$',  # Strings without format placeholders
        ]

        # Skip common safe phrases that contain sensitive keywords
        safe_phrases = [
            r"request\s+(completed|finished|started|processing)",
            r"response\s+(sent|received|processed)",
            r"data\s+(inserted|updated|deleted|saved)\s+into",
            r"(successfully|failed)\s+(request|response)",
            r"(starting|ending|completed)\s+(request|response)",
            r"no\s+(usage\s+)?data\s+found",
            r"found\s+\d+.*records",
            r"exported\s+\d+.*records",
        ]

        for pattern in non_sensitive_patterns:
            if re.search(pattern, arg_str):
                # Check if it's a safe phrase first
                for safe_pattern in safe_phrases:
                    if re.search(safe_pattern, arg_str, re.IGNORECASE):
                        return False

                # Then check if the static string mentions sensitive keywords
                if not any(
                    keyword in arg_str
                    for keyword in [
                        "request",
                        "response",
                        "data",
                        "body",
                        "payload",
                        "token",
                        "auth",
                        "credential",
                    ]
                ):
                    return False

        # Direct variable/attribute patterns that are likely sensitive
        sensitive_patterns = [
            r"\brequest\b(?!\s*(id|status|method))",  # request but not request_id, request_status, request_method
            r"\bresponse\b(?!\s*(status|code|time))",  # response but not response_status, response_code
            r"\bdata\b(?=[\.\[\s]|$)",  # data followed by . [ space or end
            r"\bbody\b(?=[\.\[\s]|$)",
            r"\bpayload\b(?=[\.\[\s]|$)",
            r"\bmessages?\b(?=[\.\[\s]|$)",
            r"\bcontent\b(?=[\.\[\s]|$)",
            r"\binput\b(?=[\.\[\s]|$)",
            r"\boutput\b(?=[\.\[\s]|$)",
            r"\bargs\b(?=[\.\[\s]|$)",
            r"\bkwargs\b(?=[\.\[\s]|$)",
            r"\bparams\b(?=[\.\[\s]|$)",
            r"\bheaders\b(?=[\.\[\s]|$)",
            r"\bapi_key\b",
            r"\btoken\b(?!\s*(name|id))",  # token but not token_name, token_id
            r"\bauth\b(?=[\.\[\s]|$)",
            r"\bcredentials?\b",
        ]

        # Check for direct variable references with context
        for pattern in sensitive_patterns:
            if re.search(pattern, arg_str):
                return True

        # Check for format strings that might interpolate sensitive data
        if self._is_format_string_with_sensitive_data(arg):
            return True

        # Check for JSON dumps or string formatting of objects
        if self._is_object_serialization(arg):
            return True

        return False

    def _is_format_string_with_sensitive_data(self, arg) -> bool:
        """Check if this is a format string that might contain sensitive data"""
        # Check for f-strings
        if isinstance(arg, ast.JoinedStr):
            for value in arg.values:
                if isinstance(value, ast.FormattedValue):
                    value_str = self._get_arg_string(value.value).lower()
                    # Check for any sensitive data patterns in f-string interpolations
                    sensitive_f_string_patterns = [
                        "request",
                        "response",
                        "data",
                        "body",
                        "content",
                        "messages",
                        "token",
                        "jwt",
                        "auth",
                        "api_key",
                        "apikey",
                        "credential",
                        "secret",
                        "password",
                        "passwd",
                    ]
                    if any(
                        pattern in value_str for pattern in sensitive_f_string_patterns
                    ):
                        return True

        # Check for .format() calls
        if isinstance(arg, ast.Call) and isinstance(arg.func, ast.Attribute):
            if arg.func.attr == "format":
                # Check the base string for suspicious patterns
                base_str = self._get_arg_string(arg.func.value).lower()
                if "{}" in base_str or "{" in base_str:
                    # Check format arguments for sensitive data
                    sensitive_format_patterns = [
                        "request",
                        "response",
                        "data",
                        "body",
                        "content",
                        "token",
                        "jwt",
                        "auth",
                        "api_key",
                        "apikey",
                        "credential",
                        "secret",
                        "password",
                        "passwd",
                    ]
                    for format_arg in arg.args:
                        format_str = self._get_arg_string(format_arg).lower()
                        if any(
                            pattern in format_str
                            for pattern in sensitive_format_patterns
                        ):
                            return True

        return False

    def _is_object_serialization(self, arg) -> bool:
        """Check if this is serializing an object that might contain sensitive data"""
        arg_str = self._get_arg_string(arg)

        # Check for json.dumps() calls
        if isinstance(arg, ast.Call):
            if (
                isinstance(arg.func, ast.Attribute)
                and arg.func.attr == "dumps"
                and isinstance(arg.func.value, ast.Name)
                and arg.func.value.id == "json"
            ):
                return True

            # Check for str() calls on potentially sensitive objects
            if (
                isinstance(arg.func, ast.Name)
                and arg.func.id == "str"
                and len(arg.args) > 0
            ):
                obj_str = self._get_arg_string(arg.args[0]).lower()
                if any(
                    pattern in obj_str
                    for pattern in ["request", "response", "data", "body"]
                ):
                    return True

        return False

    def _get_violation_reason(self, arg) -> str:
        """Get a human-readable reason for the violation"""
        arg_str = self._get_arg_string(arg).lower()

        if any(
            pattern in arg_str
            for pattern in [
                "jwt",
                "token",
                "api_key",
                "apikey",
                "auth",
                "credential",
                "secret",
                "password",
                "passwd",
            ]
        ):
            return "Potentially logging authentication/secret data (JWT, token, API key, etc.)"
        elif "request" in arg_str:
            return "Potentially logging request data"
        elif "response" in arg_str:
            return "Potentially logging response data"
        elif any(
            pattern in arg_str for pattern in ["data", "body", "payload", "content"]
        ):
            return "Potentially logging sensitive data/body/content"
        elif any(pattern in arg_str for pattern in ["messages", "input", "output"]):
            return "Potentially logging message/input/output data"
        else:
            return "Potentially logging sensitive data"

    def _get_call_string(self, node) -> str:
        """Get string representation of the function call"""
        try:
            if hasattr(ast, "unparse"):
                return ast.unparse(node)
            else:
                # Fallback for older Python versions
                return f"{self._get_arg_string(node.func)}(...)"
        except:
            return "logger.info(...)"

    def _get_arg_string(self, arg) -> str:
        """Get string representation of an argument"""
        try:
            if hasattr(ast, "unparse"):
                return ast.unparse(arg)
            else:
                # Fallback for older Python versions
                if isinstance(arg, ast.Name):
                    return arg.id
                elif isinstance(arg, ast.Attribute):
                    return f"{self._get_arg_string(arg.value)}.{arg.attr}"
                elif isinstance(arg, ast.Str):
                    return repr(arg.s)
                elif isinstance(arg, ast.Constant):
                    return repr(arg.value)
                else:
                    return str(type(arg).__name__)
        except:
            return "unknown"


def check_sensitive_logging(base_dir: str) -> List[Dict[str, Any]]:
    """
    Check for logger.info() statements that might log sensitive data.

    Args:
        base_dir: Base directory to scan (typically the litellm root)

    Returns:
        List of violations found
    """
    detector = SensitiveLogDetector()
    all_violations = []

    # Directories to scan - only main litellm codebase
    scan_dirs = ["litellm", "enterprise"]  # Include enterprise directory if it exists

    # Directories to exclude (third-party code, venvs, etc.)
    exclude_dirs = {
        "venv",
        "venv313",
        ".venv",
        "env",
        ".env",
        "node_modules",
        "__pycache__",
        ".git",
        "build",
        "dist",
        ".tox",
        "clean_env",
        "litellm_env",
        "myenv",
        "py313_env",
        "venv_sip_bypass",
        "mypyc_env",
    }

    for scan_dir in scan_dirs:
        dir_path = os.path.join(base_dir, scan_dir)
        if not os.path.exists(dir_path):
            print(f"Warning: Directory {dir_path} does not exist, skipping.")
            continue

        print(f"Scanning directory: {dir_path}")

        for root, dirs, files in os.walk(dir_path):
            # Skip excluded directories
            dirs[:] = [d for d in dirs if d not in exclude_dirs]

            # Skip if we're in a virtual environment or third-party directory
            relative_root = os.path.relpath(root, base_dir)
            if any(
                excluded in relative_root.split(os.sep) for excluded in exclude_dirs
            ):
                continue

            for file in files:
                if file.endswith(".py"):
                    file_path = os.path.join(root, file)
                    relative_path = os.path.relpath(file_path, base_dir)

                    # Skip files that are clearly third-party or generated
                    if any(excluded in relative_path for excluded in exclude_dirs):
                        continue

                    try:
                        with open(file_path, "r", encoding="utf-8") as f:
                            content = f.read()
                            tree = ast.parse(content)

                        detector.set_file(relative_path)
                        detector.visit(tree)

                    except SyntaxError as e:
                        print(f"Warning: Syntax error in file {relative_path}: {e}")
                        continue
                    except UnicodeDecodeError as e:
                        print(
                            f"Warning: Unicode decode error in file {relative_path}: {e}"
                        )
                        continue
                    except Exception as e:
                        print(f"Warning: Error processing file {relative_path}: {e}")
                        continue

    return detector.violations


def main():
    """Main function to run the sensitive logging check"""
    # Get the base directory (assume we're running from tests/code_coverage_tests/)
    ###################
    # Running locally
    ###################
    # current_dir = os.path.dirname(os.path.abspath(__file__))
    # base_dir = os.path.join(current_dir, "..", "..")
    # base_dir = os.path.abspath(base_dir)

    ###################
    # Running in CI/CD
    ###################
    base_dir = "./litellm"  # Adjust this path as needed

    print(f"Checking for sensitive logging in: {base_dir}")

    violations = check_sensitive_logging(base_dir)

    if violations:
        print(f"\n❌ Found {len(violations)} potential violations:")
        print("=" * 80)

        for i, violation in enumerate(violations, 1):
            print(f"\n{i}. {violation['file']}:{violation['line']}")
            print(f"   Reason: {violation['reason']}")
            print(f"   Call: {violation['call']}")
            print(f"   Argument: {violation['arg']}")

        print("\n" + "=" * 80)
        print("⚠️  SECURITY WARNING:")
        print("These logger.info() statements may log sensitive request/response data.")
        print("Consider changing them to logger.debug() or removing sensitive data.")
        print("This is critical for PII compliance and security.")
        print(
            "Please contact @ishaan-jaff for more details about this check. DO NOT VIOLATE THIS CHECK."
        )

        return 1  # Exit with error code
    else:
        print("\n✅ No sensitive logging violations found!")
        return 0


if __name__ == "__main__":
    exit(main())