import ast
import os

class EnterpriseImportFinder(ast.NodeVisitor):
    def __init__(self):
        self.unsafe_imports = []
        self.current_file = None
        self.in_try_block = False
        self.try_blocks = []

    def visit_Try(self, node):
        # Track that we're entering a try block
        self.in_try_block = True
        self.try_blocks.append(node)
        # Visit all nodes in the try block
        for item in node.body:
            self.visit(item)
        # Visit except blocks
        for handler in node.handlers:
            for item in handler.body:
                self.visit(item)
        # Visit else block if it exists
        for item in node.orelse:
            self.visit(item)
        # Visit finally block if it exists
        for item in node.finalbody:
            self.visit(item)
        # We're leaving the try block
        self.try_blocks.pop()
        self.in_try_block = len(self.try_blocks) > 0

    def visit_Import(self, node):
        # Check for direct imports of litellm_enterprise
        for name in node.names:
            if "litellm_enterprise" in name.name or "enterprise" in name.name:
                if not self.in_try_block:
                    self.unsafe_imports.append({
                        "file": self.current_file,
                        "line": node.lineno,
                        "import": name.name,
                        "context": "direct import"
                    })
        self.generic_visit(node)

    def visit_ImportFrom(self, node):
        # Check for from litellm_enterprise imports
        if node.module and ("litellm_enterprise" in node.module or "enterprise" in node.module):
            if not self.in_try_block:
                self.unsafe_imports.append({
                    "file": self.current_file,
                    "line": node.lineno,
                    "import": f"from {node.module}",
                    "context": "from import"
                })
        self.generic_visit(node)

def find_unsafe_enterprise_imports_in_file(file_path):
    with open(file_path, "r") as file:
        tree = ast.parse(file.read(), filename=file_path)
    finder = EnterpriseImportFinder()
    finder.current_file = file_path
    finder.visit(tree)
    return finder.unsafe_imports

def find_unsafe_enterprise_imports_in_directory(directory):
    unsafe_imports = []
    for root, _, files in os.walk(directory):
        for file in files:
            if file.endswith(".py"):
                file_path = os.path.join(root, file)
                imports = find_unsafe_enterprise_imports_in_file(file_path)
                if imports:
                    unsafe_imports.extend(imports)
    return unsafe_imports

if __name__ == "__main__":
    # Check for unsafe enterprise imports in the litellm directory
    directory_path = "./litellm"
    unsafe_imports = find_unsafe_enterprise_imports_in_directory(directory_path)
    
    if unsafe_imports:
        print("🚨 UNSAFE ENTERPRISE IMPORTS FOUND (not in try-except blocks):")
        for imp in unsafe_imports:
            print(f"File: {imp['file']}")
            print(f"Line: {imp['line']}")
            print(f"Import: {imp['import']}")
            print(f"Context: {imp['context']}")
            print("---")
        
        # Raise exception to fail CI/CD
        raise Exception(
            "🚨 Unsafe enterprise imports found. All enterprise imports must be wrapped in try-except blocks."
        )
    else:
        print("✅ No unsafe enterprise imports found.")