"""
LangGraph adapter for UASS integration.
"""

from typing import Optional, Dict, Any
from ..client import UASSClient
from ..adapters.base import FrameworkAdapter


class LangGraphAdapter(FrameworkAdapter):
    """
    Adapter for integrating LangGraph with UASS.
    
    Usage:
        from uass import UASSClient, RedisBackend, LangGraphAdapter
        
        client = UASSClient(backend=RedisBackend())
        adapter = LangGraphAdapter(client)
        
        # In LangGraph node
        adapter.sync_to_uass("session_123", "agent_1", graph_state)
        
        # Read state from another framework
        state = adapter.sync_from_uass("session_123", "agent_1")
    """
    
    def sync_to_uass(
        self,
        session_id: str,
        agent_id: str,
        graph_state: Dict[str, Any],
    ) -> bool:
        """
        Sync LangGraph state to UASS.
        
        Args:
            session_id: Session identifier
            agent_id: Agent identifier
            graph_state: LangGraph state dictionary
            
        Returns:
            True if successful, False otherwise
        """
        # Extract common state from LangGraph format
        variables = graph_state.get("variables", {})
        messages = graph_state.get("messages", [])
        
        # Convert LangGraph messages to UASS format
        uass_messages = []
        for msg in messages:
            if isinstance(msg, dict):
                uass_messages.append(msg)
            else:
                # Handle LangChain message objects
                uass_messages.append({
                    "type": getattr(msg, "type", "unknown"),
                    "content": getattr(msg, "content", str(msg)),
                })
        
        return self.client.write_state(
            session_id=session_id,
            agent_id=agent_id,
            framework="langgraph",
            variables=variables,
            messages=uass_messages,
            framework_state=graph_state,
        )
    
    def sync_from_uass(
        self,
        session_id: str,
        agent_id: str,
    ) -> Optional[Dict[str, Any]]:
        """
        Sync state from UASS to LangGraph format.
        
        Args:
            session_id: Session identifier
            agent_id: Agent identifier
            
        Returns:
            LangGraph state dictionary, or None if not found
        """
        unified_state = self.client.read_full_state(session_id, agent_id)
        if unified_state is None:
            return None
        
        # Reconstruct LangGraph state format
        # Exclude messages from framework_state since we use the converted ones
        framework_state = unified_state.state.framework_state.copy()
        framework_state.pop("messages", None)  # Remove messages if present
        
        langgraph_state = {
            "variables": unified_state.state.variables,
            "messages": unified_state.state.messages,  # Use converted messages
            **framework_state,  # Preserve other LangGraph-specific state
        }
        
        return langgraph_state
    
    def create_checkpoint_hook(self, session_id: str, agent_id: str):
        """
        Create a checkpoint hook for LangGraph.
        
        Returns a function that can be used as a checkpoint callback.
        """
        def checkpoint_hook(state: Dict[str, Any]):
            """Hook function for LangGraph checkpoints."""
            self.sync_to_uass(session_id, agent_id, state)
        
        return checkpoint_hook

