"""
Haystack adapter for UASS integration.
"""

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


class HaystackAdapter(FrameworkAdapter):
    """
    Adapter for integrating Haystack with UASS.
    
    Usage:
        from uass import UASSClient, RedisBackend, HaystackAdapter
        
        client = UASSClient(backend=RedisBackend())
        adapter = HaystackAdapter(client)
        
        # Sync Haystack agent state
        adapter.sync_to_uass("session_123", "agent_1", agent, messages, metadata)
        
        # 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,
        agent: Any = None,
        messages: Optional[List[Any]] = None,
        metadata: Optional[Dict[str, Any]] = None,
        variables: Optional[Dict[str, Any]] = None,
    ) -> bool:
        """
        Sync Haystack state to UASS.
        
        Args:
            session_id: Session identifier
            agent_id: Agent identifier
            agent: Haystack agent instance (optional)
            messages: Haystack messages (optional)
            metadata: Additional metadata (optional)
            variables: Additional variables (optional)
            
        Returns:
            True if successful, False otherwise
        """
        uass_variables = variables or {}
        context = {}
        uass_messages = []
        actions = []
        
        if agent:
            # Extract agent information
            uass_variables["agent_name"] = getattr(agent, "name", getattr(agent, "agent_name", agent_id))
            if hasattr(agent, "description"):
                uass_variables["agent_description"] = getattr(agent, "description", "")
            if hasattr(agent, "prompt_template"):
                context["prompt_template"] = str(getattr(agent, "prompt_template", ""))
        
        if metadata:
            # Store metadata in context
            context.update(metadata)
        
        if messages:
            # Convert Haystack messages to UASS format
            for msg in messages:
                if isinstance(msg, dict):
                    uass_messages.append(msg)
                else:
                    # Handle Haystack message objects
                    uass_messages.append({
                        "role": getattr(msg, "role", getattr(msg, "from_", "unknown")),
                        "content": getattr(msg, "content", getattr(msg, "message", str(msg))),
                        "meta": getattr(msg, "meta", {}),
                    })
        
        framework_state = {}
        if agent:
            framework_state["agent_type"] = type(agent).__name__
        
        return self.client.write_state(
            session_id=session_id,
            agent_id=agent_id,
            framework="haystack",
            variables=uass_variables,
            context=context,
            messages=uass_messages,
            actions=actions,
            framework_state=framework_state,
        )
    
    def sync_from_uass(
        self,
        session_id: str,
        agent_id: str,
    ) -> Optional[Dict[str, Any]]:
        """
        Sync state from UASS to Haystack format.
        
        Args:
            session_id: Session identifier
            agent_id: Agent identifier
            
        Returns:
            Dictionary with Haystack-compatible state, or None if not found
        """
        unified_state = self.client.read_full_state(session_id, agent_id)
        if unified_state is None:
            return None
        
        # Return Haystack-compatible format
        return {
            "messages": unified_state.state.messages,
            "variables": unified_state.state.variables,
            "context": unified_state.state.context,
            "actions": unified_state.state.actions,
            "framework_state": unified_state.state.framework_state,
        }

