import ast
from typing import Dict, Set, Tuple, List


class ClassMethodMap:
    """Stores method name mappings for all classes in the code."""
    def __init__(self):
        # Maps: original_class_name -> {original_method_name -> obfuscated_method_name}
        self.class_methods: Dict[str, Dict[str, str]] = {}
        
        # Maps: original_class_name -> {original_attr_name -> obfuscated_attr_name}
        self.class_attributes: Dict[str, Dict[str, str]] = {}
        
        # Maps: original_class_name -> obfuscated_class_name
        self.class_renames: Dict[str, str] = {}
        
        # Track inheritance relationships: child_class -> [parent_classes]
        self.inheritance: Dict[str, List[str]] = {}


class ClassAnalyzer(ast.NodeVisitor):
    """
    Pre-analyzes classes to ensure consistent renaming of methods and attributes.
    This is crucial for making self.method() calls match def method() definitions.
    """
    def __init__(self, name_generator):
        self.name_generator = name_generator
        self.method_map = ClassMethodMap()
        self.current_class = None
        
        # To avoid duplicate scanning
        self.scanned_classes: Set[str] = set()
        
        # Track method calls within each class
        self.method_calls: Dict[str, Set[str]] = {}

    def analyze(self, tree: ast.AST) -> ClassMethodMap:
        """Analyzes the entire AST and returns populated method mappings."""
        self.visit(tree)
        self._resolve_inheritance()
        self._ensure_consistent_method_mapping()
        return self.method_map
    
    def visit_ClassDef(self, node: ast.ClassDef):
        """Process a class definition and map its methods."""
        prev_class = self.current_class
        self.current_class = node.name
        
        # Skip if already processed this class
        if node.name in self.scanned_classes:
            self.current_class = prev_class
            return
        
        # Initialize method calls tracking for this class
        self.method_calls[node.name] = set()
        
        # Record class inheritance
        parent_classes = []
        for base in node.bases:
            if isinstance(base, ast.Name):
                parent_classes.append(base.id)
        
        if parent_classes:
            self.method_map.inheritance[node.name] = parent_classes
        
        # Initialize mappings for this class
        if node.name not in self.method_map.class_methods:
            self.method_map.class_methods[node.name] = {}
        
        if node.name not in self.method_map.class_attributes:
            self.method_map.class_attributes[node.name] = {}
        
        # Create a consistent obfuscated name for this class
        if node.name not in self.method_map.class_renames:
            new_name = self.name_generator.generate_name()
            self.method_map.class_renames[node.name] = new_name
        
        # Process all method definitions in the class
        for item in node.body:
            # Methods
            if isinstance(item, ast.FunctionDef):
                # Skip dunder methods
                if not (item.name.startswith('__') and item.name.endswith('__')):
                    # Generate a consistent obfuscated name for this method
                    new_name = self.name_generator.generate_name()
                    self.method_map.class_methods[node.name][item.name] = new_name
                    
                # Visit the method body to find self.method() calls
                self.visit(item)
            
            # Attributes in assignments
            elif isinstance(item, ast.Assign):
                self.visit_attribute_assign(item)
            else:
                # Visit other nodes (like if statements that might contain self.method calls)
                self.visit(item)
        
        self.scanned_classes.add(node.name)
        
        # Visit any nested classes
        for item in node.body:
            if isinstance(item, ast.ClassDef):
                self.visit(item)
        
        self.current_class = prev_class
    
    def visit_attribute_assign(self, node):
        """Process attribute assignments like self.attr = value"""
        if not self.current_class:
            return
            
        for target in node.targets:
            if isinstance(target, ast.Attribute) and isinstance(target.value, ast.Name):
                if target.value.id == 'self':
                    # This is a self.attribute assignment
                    attr_name = target.attr
                    if attr_name not in self.method_map.class_attributes[self.current_class]:
                        new_name = self.name_generator.generate_name()
                        self.method_map.class_attributes[self.current_class][attr_name] = new_name
        
        # Visit the value part of the assignment for nested self.method() calls
        self.visit(node.value)
    
    def visit_Attribute(self, node):
        """Track self.method references to ensure consistent naming"""
        if self.current_class:
            is_self_method, method_name = get_method_name(node)
            if is_self_method:
                # Record this method call for consistency checks later
                self.method_calls[self.current_class].add(method_name)
        
        # Continue traversing
        self.generic_visit(node)
    
    def _ensure_consistent_method_mapping(self):
        """
        Make sure that all methods called via self.method() have a mapping,
        even if they're not defined in the class.
        """
        for class_name, method_calls in self.method_calls.items():
            if class_name not in self.method_map.class_methods:
                continue
                
            class_methods = self.method_map.class_methods[class_name]
            for method_name in method_calls:
                if method_name not in class_methods:
                    # Skip dunder methods
                    if method_name.startswith('__') and method_name.endswith('__'):
                        continue
                    # Existing check: mapping is generated only once.
                    new_name = self.name_generator.generate_name()
                    class_methods[method_name] = new_name
    
    def _resolve_inheritance(self):
        """
        Ensure child classes inherit method mappings from parent classes.
        This ensures that overridden methods use the same obfuscated name.
        """
        # Process inheritance depth-first to handle multi-level inheritance
        def process_inheritance(class_name):
            if class_name not in self.method_map.inheritance:
                return
            
            for parent in self.method_map.inheritance[class_name]:
                # Process parent's inheritance first
                process_inheritance(parent)
                
                # Skip if parent isn't in our mappings (external class)
                if parent not in self.method_map.class_methods:
                    continue
                
                # Inherit parent's methods if not overridden
                for method_name, obf_name in self.method_map.class_methods[parent].items():
                    if method_name not in self.method_map.class_methods[class_name]:
                        self.method_map.class_methods[class_name][method_name] = obf_name
        
        # Process inheritance for each class
        for class_name in list(self.method_map.class_methods.keys()):
            process_inheritance(class_name)


def get_method_name(node: ast.Attribute) -> Tuple[bool, str]:
    """
    Helper function to determine if an attribute is a self.method() call.
    Returns (is_self_method, method_name)
    """
    if isinstance(node.value, ast.Name) and node.value.id == 'self':
        return True, node.attr
    return False, ""


def update_obfuscator_with_class_mappings(obfuscator, class_map: ClassMethodMap):
    """
    Updates the main obfuscator with class method and attribute mappings
    to ensure consistent renaming across the codebase.
    """
    # Update class name mappings in global_var_renames
    for orig_name, obf_name in class_map.class_renames.items():
        obfuscator.global_var_renames[orig_name] = obf_name
    
    # Update class attr mapping with our analyzed data
    for class_name, class_obf_name in class_map.class_renames.items():
        # Initialize if needed
        if class_obf_name not in obfuscator.class_attr_mapping:
            obfuscator.class_attr_mapping[class_obf_name] = {}
        
        # Copy method mappings
        if class_name in class_map.class_methods:
            for method, obf_method in class_map.class_methods[class_name].items():
                obfuscator.class_attr_mapping[class_obf_name][method] = obf_method
        
        # Copy attribute mappings
        if class_name in class_map.class_attributes:
            for attr, obf_attr in class_map.class_attributes[class_name].items():
                obfuscator.class_attr_mapping[class_obf_name][attr] = obf_attr