"""
transpiler/visitor.py — Unified visitor and transformer hierarchy for RPy IR.
"""

from __future__ import annotations
from typing import Any, List, Optional, TypeVar, Generic, Union, cast
from transpiler.middleend import ir

T = TypeVar("T")

from transpiler.infrastructure import logger

class BaseIRVisitor:
    """Base class for all IR traversals. 
    Dispatches to visit_ClassName methods.
    """
    def visit(self, node: Any) -> Any:
        if node is None:
            return None
        if isinstance(node, list):
            return [self.visit(item) for item in node if item is not None]
        
        method_name = f"visit_{node.__class__.__name__}"
        visitor_fn = getattr(self, method_name, None)
        if visitor_fn is None:
            logger.debug_compiler("Visitor", f"Missing visitor {method_name} in {self.__class__.__name__}")
            visitor_fn = self.generic_visit
        return visitor_fn(node)

    def generic_visit(self, node: ir.IRNode) -> Any:
        """Default behavior: recurse into all children."""
        if hasattr(node, "__dict__"):
            for attr, value in node.__dict__.items():
                if isinstance(value, ir.IRNode):
                    self.visit(value)
                elif isinstance(value, list):
                    for item in value:
                        if isinstance(item, ir.IRNode):
                            self.visit(item)
                elif isinstance(value, ir.IRIfElseExpression):
                    self.visit(value.condition)
                    self.visit(value.then_expr)
                    self.visit(value.else_expr)
        return node

class IRVisitor(BaseIRVisitor):
    """Read-only visitor for analysis passes."""
    def generic_visit(self, node: ir.IRNode) -> None:
        super().generic_visit(node)
        return None

class IRTransformer(BaseIRVisitor):
    """Mutating visitor for transformation passes.
    Returns a new or modified node.
    """
    def visit(self, node: Any) -> Any:
        if node is None:
            return None
        if isinstance(node, list):
            new_list = []
            for item in node:
                res = self.visit(item)
                if res is None:
                    continue # Deletion
                if isinstance(res, list):
                    new_list.extend(res) # Splicing
                else:
                    new_list.append(res)
            return new_list
        
        method_name = f"visit_{node.__class__.__name__}"
        visitor_fn = getattr(self, method_name, None)
        if visitor_fn is None:
            if method_name == "visit_IRIfElseExpression":
                import sys
                print(f"DEBUG: No transformer for IfElseExpression in {self.__class__.__name__}", file=sys.stderr)
            visitor_fn = self.generic_visit
        return visitor_fn(node)

    def generic_visit(self, node: ir.IRNode) -> ir.IRNode:
        """Mutates the node in-place by visiting children and updating attributes."""
        if hasattr(node, "__dict__"):
            for attr, value in node.__dict__.items():
                if isinstance(value, ir.IRNode):
                    new_val = self.visit(value)
                    setattr(node, attr, new_val)
                elif isinstance(value, list):
                    new_list = self.visit(value)
                    setattr(node, attr, new_list)
                elif isinstance(value, ir.IRIfElseExpression):
                    value.condition = self.visit(value.condition)
                    value.then_expr = self.visit(value.then_expr)
                    value.else_expr = self.visit(value.else_expr)
        return node

def copy_metadata(old: ir.IRNode, new: ir.IRNode) -> ir.IRNode:
    """Helper to preserve source information when replacing nodes."""
    if hasattr(old, "lineno") and hasattr(new, "lineno"):
        new.lineno = old.lineno
    if hasattr(old, "col_offset") and hasattr(new, "col_offset"):
        new.col_offset = old.col_offset
    return new
