"""
Main UASS client for interacting with unified state.
"""

from typing import Optional, Dict, Any, List
from datetime import datetime

from .backends import StateBackend
from .schemas import (
    UnifiedState,
    AgentState,
    StateMetadata,
    TokenUsage,
    FrameworkType,
)


class UASSClient:
    """
    Main client for unified agent state management.
    
    This client provides a simple interface for any framework to read/write
    state in a unified format.
    """
    
    def __init__(self, backend: StateBackend):
        """
        Initialize UASS client.
        
        Args:
            backend: State storage backend (RedisBackend, MemoryBackend, etc.)
        """
        self.backend = backend
    
    def write_state(
        self,
        session_id: str,
        agent_id: str,
        framework: str,
        state: Optional[Dict[str, Any]] = None,
        variables: Optional[Dict[str, Any]] = None,
        context: Optional[Dict[str, Any]] = None,
        memory: Optional[Dict[str, Any]] = None,
        tokens: Optional[TokenUsage] = None,
        messages: Optional[List[Dict[str, Any]]] = None,
        actions: Optional[List[Dict[str, Any]]] = None,
        framework_state: Optional[Dict[str, Any]] = None,
        framework_version: Optional[str] = None,
    ) -> bool:
        """
        Write state to unified store.
        
        Args:
            session_id: Unique session identifier
            agent_id: Agent identifier within session
            framework: Framework name (langgraph, autogen, crewai, etc.)
            state: Complete state dict (alternative to individual params)
            variables: Agent variables
            context: Agent context
            memory: Agent memory
            tokens: Token usage information
            messages: Conversation messages
            actions: Agent actions
            framework_state: Framework-specific opaque state
            framework_version: Framework version string
            
        Returns:
            True if successful, False otherwise
        """
        # Read existing state to preserve metadata
        existing = self.backend.read_state(session_id, agent_id)
        
        now = datetime.utcnow()
        
        if existing:
            # Update existing state
            metadata = StateMetadata(
                created_at=existing.metadata.created_at,
                updated_at=now,
                version=existing.metadata.version + 1,
                framework=framework,
                framework_version=framework_version,
            )
            agent_state = existing.state
        else:
            # Create new state
            metadata = StateMetadata(
                created_at=now,
                updated_at=now,
                version=1,
                framework=framework,
                framework_version=framework_version,
            )
            agent_state = AgentState()
        
        # Update state fields
        if state:
            # If complete state dict provided, use it
            agent_state = AgentState.from_dict(state)
        else:
            # Update individual fields
            if variables is not None:
                agent_state.variables.update(variables)
            if context is not None:
                agent_state.context.update(context)
            if memory is not None:
                agent_state.memory.update(memory)
            if tokens is not None:
                agent_state.tokens = tokens
            if messages is not None:
                agent_state.messages = messages
            if actions is not None:
                agent_state.actions = actions
            if framework_state is not None:
                agent_state.framework_state.update(framework_state)
        
        unified_state = UnifiedState(
            session_id=session_id,
            agent_id=agent_id,
            state=agent_state,
            metadata=metadata,
        )
        
        return self.backend.write_state(unified_state)
    
    def read_state(
        self,
        session_id: str,
        agent_id: str,
    ) -> Optional[Dict[str, Any]]:
        """
        Read state from unified store.
        
        Args:
            session_id: Session identifier
            agent_id: Agent identifier
            
        Returns:
            State dictionary if found, None otherwise
        """
        unified_state = self.backend.read_state(session_id, agent_id)
        if unified_state is None:
            return None
        
        return unified_state.state.to_dict()
    
    def read_full_state(
        self,
        session_id: str,
        agent_id: str,
    ) -> Optional[UnifiedState]:
        """
        Read full unified state (including metadata).
        
        Args:
            session_id: Session identifier
            agent_id: Agent identifier
            
        Returns:
            UnifiedState if found, None otherwise
        """
        return self.backend.read_state(session_id, agent_id)
    
    def delete_state(self, session_id: str, agent_id: str) -> bool:
        """
        Delete state from unified store.
        
        Args:
            session_id: Session identifier
            agent_id: Agent identifier
            
        Returns:
            True if successful, False otherwise
        """
        return self.backend.delete_state(session_id, agent_id)
    
    def list_agents(self, session_id: str) -> List[str]:
        """
        List all agents in a session.
        
        Args:
            session_id: Session identifier
            
        Returns:
            List of agent IDs
        """
        return self.backend.list_agents(session_id)
    
    def write_shared_state(
        self,
        session_id: str,
        state: Dict[str, Any],
        framework: str = "shared",
    ) -> bool:
        """
        Write shared state accessible by all agents in a session.
        
        Args:
            session_id: Session identifier
            state: State dictionary
            framework: Framework identifier (default: "shared")
            
        Returns:
            True if successful, False otherwise
        """
        now = datetime.utcnow()
        agent_state = AgentState.from_dict(state)
        
        metadata = StateMetadata(
            created_at=now,
            updated_at=now,
            version=1,
            framework=framework,
        )
        
        unified_state = UnifiedState(
            session_id=session_id,
            agent_id="shared",
            state=agent_state,
            metadata=metadata,
        )
        
        return self.backend.write_shared_state(session_id, unified_state)
    
    def read_shared_state(self, session_id: str) -> Optional[Dict[str, Any]]:
        """
        Read shared state for a session.
        
        Args:
            session_id: Session identifier
            
        Returns:
            State dictionary if found, None otherwise
        """
        unified_state = self.backend.read_shared_state(session_id)
        if unified_state is None:
            return None
        
        return unified_state.state.to_dict()
    
    def update_tokens(
        self,
        session_id: str,
        agent_id: str,
        prompt_tokens: int = 0,
        completion_tokens: int = 0,
        cost_usd: float = 0.0,
    ) -> bool:
        """
        Update token usage for an agent.
        
        Args:
            session_id: Session identifier
            agent_id: Agent identifier
            prompt_tokens: Number of prompt tokens
            completion_tokens: Number of completion tokens
            cost_usd: Cost in USD
            
        Returns:
            True if successful, False otherwise
        """
        existing = self.backend.read_state(session_id, agent_id)
        if existing is None:
            return False
        
        tokens = TokenUsage(
            prompt_tokens=prompt_tokens,
            completion_tokens=completion_tokens,
            total_tokens=prompt_tokens + completion_tokens,
            cost_usd=cost_usd,
        )
        
        existing.state.tokens = tokens
        
        unified_state = UnifiedState(
            session_id=session_id,
            agent_id=agent_id,
            state=existing.state,
            metadata=StateMetadata(
                created_at=existing.metadata.created_at,
                updated_at=datetime.utcnow(),
                version=existing.metadata.version + 1,
                framework=existing.metadata.framework,
                framework_version=existing.metadata.framework_version,
            ),
        )
        
        return self.backend.write_state(unified_state)

