import ast
from typing import Dict, Optional

class AttributeTransformer(ast.NodeTransformer):
    """
    Transforms attribute access expressions (obj.attr) consistently,
    especially for class method calls and attribute accesses.
    """
    
    def __init__(self, class_attr_mapping: Dict[str, Dict[str, str]], class_renames: Dict[str, str]):
        """
        Initialize with mapping dictionaries.
        
        Args:
            class_attr_mapping: Maps class_name -> {attr_name -> obfuscated_attr_name}
            class_renames: Maps original_class_name -> obfuscated_class_name
        """
        self.class_attr_mapping = class_attr_mapping
        self.class_renames = class_renames
        self.current_class: Optional[str] = None
    
    def visit_ClassDef(self, node):
        """Keep track of the current class being processed"""
        old_class = self.current_class
        
        # Get the obfuscated name for this class
        class_name = node.name
        if class_name in self.class_renames:
            self.current_class = self.class_renames[class_name]
        else:
            self.current_class = class_name
        
        # Process the class body
        node = self.generic_visit(node)
        
        # Restore the previous class context
        self.current_class = old_class
        return node
    
    def visit_Attribute(self, node):
        """Transform attribute access like self.method to use consistent names"""
        # First process any nested attributes
        node = self.generic_visit(node)
        
        # Handle self.attr references within a class
        if isinstance(node.value, ast.Name) and node.value.id == 'self' and self.current_class:
            if self.current_class in self.class_attr_mapping:
                attr_map = self.class_attr_mapping[self.current_class]
                # Only substitute if node.attr is still in its original form (a mapping key)
                if node.attr in attr_map and node.attr != attr_map[node.attr]:
                    node.attr = attr_map[node.attr]
        
        return node