"""
LangGraph StateGraph Auto-Patching for Zero-Config Tracing

This module provides automatic tracing for all LangGraph StateGraph instances
without requiring any code changes in agent files.

How it works:
1. patch_langgraph() patches StateGraph.__init__, add_node, and compile
2. When agents create StateGraphs, they're automatically traced
3. Nested agent calls are detected and create proper hierarchy

v0.6.0
"""

import functools
import logging
import re
import threading
from typing import Type, Optional, Any, Callable

logger = logging.getLogger(__name__)

# Module state
_patched = False
_patch_lock = threading.Lock()
_original_methods = {}

# Minimum supported LangGraph version
MIN_LANGGRAPH_VERSION = "0.2.0"


def _infer_trace_name(state_class: Type) -> str:
    """
    Infer trace name from state class name.
    
    ProductivityAgentState -> "productivity_agent"
    GlobalSupervisorState -> "global_supervisor"
    MyAgentState -> "my_agent"
    
    Algorithm:
    1. Get class name (e.g., "ProductivityAgentState")
    2. Strip "State" suffix if present
    3. Convert CamelCase to snake_case
    """
    if state_class is None:
        return "unknown"
    
    name = state_class.__name__
    
    # Strip "State" suffix
    if name.endswith("State"):
        name = name[:-5]  # Remove "State"
    
    # Convert CamelCase to snake_case
    # Insert underscore before uppercase letters (except first)
    snake_case = re.sub(r'(?<!^)(?=[A-Z])', '_', name).lower()
    
    return snake_case


def _check_langgraph_version() -> bool:
    """
    Check if LangGraph version is compatible.
    
    Returns True if LangGraph >= MIN_LANGGRAPH_VERSION, False otherwise.
    """
    try:
        import langgraph
        version = getattr(langgraph, '__version__', '0.0.0')
        
        # Parse version (handle formats like "0.2.0", "0.2.0rc1", etc.)
        version_parts = re.match(r'^(\d+)\.(\d+)\.(\d+)', version)
        if not version_parts:
            logger.warning(f"Could not parse LangGraph version: {version}")
            return True  # Assume compatible if can't parse
        
        major, minor, patch = map(int, version_parts.groups())
        min_parts = MIN_LANGGRAPH_VERSION.split('.')
        min_major, min_minor, min_patch = map(int, min_parts)
        
        # Compare versions
        if major > min_major:
            return True
        if major == min_major and minor > min_minor:
            return True
        if major == min_major and minor == min_minor and patch >= min_patch:
            return True
        
        logger.warning(
            f"LangGraph version {version} is below minimum {MIN_LANGGRAPH_VERSION}. "
            "Auto-patching may not work correctly."
        )
        return True  # Still return True to attempt patching
        
    except ImportError:
        logger.debug("LangGraph not installed, cannot check version")
        return False
    except Exception as e:
        logger.debug(f"Error checking LangGraph version: {e}")
        return True  # Assume compatible on error


def patch_langgraph() -> bool:
    """
    Apply automatic tracing patches to LangGraph StateGraph.
    
    Thread-safe, idempotent (calling twice returns False, not error).
    
    Patches:
    - StateGraph.__init__: Infers trace name from state class
    - StateGraph.add_node: Wraps functions with @traced
    - StateGraph.compile: Wraps ainvoke to create traces
    
    Returns:
        True if patching succeeded (first call)
        False if already patched or failed
    """
    global _patched, _original_methods
    
    with _patch_lock:
        if _patched:
            logger.debug("LangGraph already patched, skipping")
            return False
        
        try:
            # Check version compatibility
            if not _check_langgraph_version():
                logger.warning("LangGraph not available or incompatible")
                return False
            
            # Import LangGraph
            from langgraph.graph import StateGraph
            
            # Store original methods
            _original_methods['__init__'] = StateGraph.__init__
            _original_methods['add_node'] = StateGraph.add_node
            _original_methods['compile'] = StateGraph.compile
            
            # Import tracing components
            from .decorator import traced
            from .graph_wrapper import _handle_nested_agent, _handle_top_level_agent, _handle_hitl_continuation
            from .context import get_current_trace, get_current_observation
            from .redis_publisher import is_tracing_initialized, get_redis_client, get_trace_for_thread
            
            # ================================================================
            # PATCH 1: __init__ - Store inferred trace name
            # ================================================================
            original_init = _original_methods['__init__']
            
            @functools.wraps(original_init)
            def patched_init(self, state_schema=None, config_schema=None):
                try:
                    # Call original init
                    original_init(self, state_schema, config_schema)
                    
                    # Infer and store trace name
                    if state_schema is not None:
                        self._swisper_trace_name = _infer_trace_name(state_schema)
                    else:
                        self._swisper_trace_name = "graph"
                        
                except Exception as e:
                    logger.debug(f"Error in patched __init__: {e}")
                    # Ensure original init is called even on error
                    try:
                        original_init(self, state_schema, config_schema)
                    except Exception:
                        pass
                    raise
            
            # ================================================================
            # PATCH 2: add_node - Auto-wrap with @traced
            # ================================================================
            original_add_node = _original_methods['add_node']
            
            @functools.wraps(original_add_node)
            def patched_add_node(self, node: str, action: Callable, **kwargs):
                try:
                    # Wrap action with @traced decorator
                    wrapped_action = traced(
                        name=node,
                        observation_type="AUTO"
                    )(action)
                    
                    return original_add_node(self, node, wrapped_action, **kwargs)
                    
                except Exception as e:
                    logger.debug(f"Failed to wrap node {node}: {e}, using original")
                    return original_add_node(self, node, action, **kwargs)
            
            # ================================================================
            # PATCH 3: compile - Wrap ainvoke for trace creation
            # ================================================================
            original_compile = _original_methods['compile']
            
            @functools.wraps(original_compile)
            def patched_compile(self, *args, **kwargs):
                try:
                    compiled = original_compile(self, *args, **kwargs)
                    original_ainvoke = compiled.ainvoke
                    
                    # Get trace name (set by __init__ or set_trace_name)
                    trace_name = getattr(self, '_swisper_trace_name', 'graph')
                    
                    @functools.wraps(original_ainvoke)
                    async def traced_ainvoke(input_state, config=None, **invoke_kwargs):
                        """Wrapped ainvoke that creates/continues traces"""
                        
                        # Pre-flight checks
                        if not is_tracing_initialized():
                            return await original_ainvoke(input_state, config, **invoke_kwargs)
                        
                        if not get_redis_client():
                            return await original_ainvoke(input_state, config, **invoke_kwargs)
                        
                        # Check for existing trace (nested agent)
                        existing_trace_id = get_current_trace()
                        existing_parent_obs = get_current_observation()
                        
                        if existing_trace_id:
                            # Nested agent call
                            return await _handle_nested_agent(
                                original_ainvoke, input_state, config, invoke_kwargs,
                                trace_name, existing_trace_id, existing_parent_obs
                            )
                        else:
                            # Top-level agent
                            return await _handle_top_level_agent(
                                original_ainvoke, input_state, config, invoke_kwargs,
                                trace_name
                            )
                    
                    compiled.ainvoke = traced_ainvoke
                    return compiled
                    
                except Exception as e:
                    logger.debug(f"Failed to wrap compile: {e}")
                    return original_compile(self, *args, **kwargs)
            
            # Apply patches
            StateGraph.__init__ = patched_init
            StateGraph.add_node = patched_add_node
            StateGraph.compile = patched_compile
            
            _patched = True
            logger.info("✅ LangGraph StateGraph patched for auto-tracing")
            return True
            
        except ImportError as e:
            logger.debug(f"LangGraph not available: {e}")
            return False
        except Exception as e:
            logger.warning(f"Failed to patch LangGraph: {e}")
            return False


def unpatch_langgraph() -> bool:
    """
    Remove patches from LangGraph StateGraph.
    
    Useful for testing cleanup.
    
    Returns:
        True if unpatching succeeded, False otherwise
    """
    global _patched, _original_methods
    
    with _patch_lock:
        if not _patched:
            return False
        
        try:
            from langgraph.graph import StateGraph
            
            if '__init__' in _original_methods:
                StateGraph.__init__ = _original_methods['__init__']
            if 'add_node' in _original_methods:
                StateGraph.add_node = _original_methods['add_node']
            if 'compile' in _original_methods:
                StateGraph.compile = _original_methods['compile']
            
            _original_methods.clear()
            _patched = False
            logger.debug("LangGraph StateGraph unpatched")
            return True
            
        except Exception as e:
            logger.debug(f"Failed to unpatch LangGraph: {e}")
            return False


def set_trace_name(graph: Any, name: str) -> None:
    """
    Explicitly set trace name for a graph (overrides inference).
    
    Use when automatic naming isn't appropriate.
    
    Args:
        graph: StateGraph instance
        name: Trace name to use
        
    Example:
        graph = StateGraph(MyCustomState)
        set_trace_name(graph, "custom_agent")
    """
    graph._swisper_trace_name = name


def is_patched() -> bool:
    """
    Check if LangGraph is currently patched.
    
    Returns:
        True if patched, False otherwise
    """
    return _patched
