"""
@file transformers/attribute_transformer.py
@brief Attribute access transformer for consistent class member renaming.
@details Ensures references like self.method and self.attr are kept consistent
         with class and attribute mappings produced by analysis.
"""

import ast
from typing import Dict, Optional

class AttributeTransformer(ast.NodeTransformer):
    """
    @brief Transform attribute access expressions for consistency.
    @details Handles self.method and self.attr translations based on provided
             mappings to keep calls and attributes aligned with obfuscated names.
    """
    
    def __init__(self, class_attr_mapping: Dict[str, Dict[str, str]], class_renames: Dict[str, str]):
        """
        @brief Initialize transformer with mapping dictionaries.
        @param class_attr_mapping Maps class_name -> {attr_name -> obfuscated_attr_name}
        @param 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):
        """
        @brief Track the current class and process its body.
        @param node ast.ClassDef being visited.
        @return ast.ClassDef Potentially modified node.
        """
        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):
        """
        @brief Transform attribute access for consistency within a class.
        @param node ast.Attribute node to transform.
        @return ast.AST Updated attribute node.
        """
        # 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