"""
Unified state schemas for cross-framework agent communication.
"""

from typing import Dict, Any, Optional, List
from dataclasses import dataclass, field
from datetime import datetime
from enum import Enum
import json


class FrameworkType(str, Enum):
    """Supported AI agent frameworks."""
    LANGGRAPH = "langgraph"
    AUTOGEN = "autogen"
    CREWAI = "crewai"
    SEMANTIC_KERNEL = "semantic_kernel"
    HAYSTACK = "haystack"
    LLAMAINDEX = "llamaindex"
    CUSTOM = "custom"


@dataclass
class TokenUsage:
    """Token usage tracking."""
    prompt_tokens: int = 0
    completion_tokens: int = 0
    total_tokens: int = 0
    cost_usd: float = 0.0
    
    def to_dict(self) -> Dict[str, Any]:
        return {
            "prompt_tokens": self.prompt_tokens,
            "completion_tokens": self.completion_tokens,
            "total_tokens": self.total_tokens,
            "cost_usd": self.cost_usd,
        }


@dataclass
class AgentState:
    """Framework-agnostic agent state."""
    # Core state
    variables: Dict[str, Any] = field(default_factory=dict)
    context: Dict[str, Any] = field(default_factory=dict)
    memory: Dict[str, Any] = field(default_factory=dict)
    
    # Token tracking
    tokens: Optional[TokenUsage] = None
    
    # Framework-specific state (opaque to UASS)
    framework_state: Dict[str, Any] = field(default_factory=dict)
    
    # History/Conversation
    messages: List[Dict[str, Any]] = field(default_factory=list)
    actions: List[Dict[str, Any]] = field(default_factory=list)
    
    def to_dict(self) -> Dict[str, Any]:
        """Convert to dictionary for serialization."""
        result = {
            "variables": self.variables,
            "context": self.context,
            "memory": self.memory,
            "framework_state": self.framework_state,
            "messages": self.messages,
            "actions": self.actions,
        }
        if self.tokens:
            result["tokens"] = self.tokens.to_dict()
        return result
    
    @classmethod
    def from_dict(cls, data: Dict[str, Any]) -> "AgentState":
        """Create from dictionary."""
        tokens = None
        if "tokens" in data and data["tokens"]:
            tokens = TokenUsage(**data["tokens"])
        
        return cls(
            variables=data.get("variables", {}),
            context=data.get("context", {}),
            memory=data.get("memory", {}),
            tokens=tokens,
            framework_state=data.get("framework_state", {}),
            messages=data.get("messages", []),
            actions=data.get("actions", []),
        )


@dataclass
class StateMetadata:
    """Metadata about the state."""
    created_at: datetime
    updated_at: datetime
    version: int = 1
    framework: str = ""
    framework_version: Optional[str] = None
    
    def to_dict(self) -> Dict[str, Any]:
        return {
            "created_at": self.created_at.isoformat(),
            "updated_at": self.updated_at.isoformat(),
            "version": self.version,
            "framework": self.framework,
            "framework_version": self.framework_version,
        }
    
    @classmethod
    def from_dict(cls, data: Dict[str, Any]) -> "StateMetadata":
        return cls(
            created_at=datetime.fromisoformat(data["created_at"]),
            updated_at=datetime.fromisoformat(data["updated_at"]),
            version=data.get("version", 1),
            framework=data.get("framework", ""),
            framework_version=data.get("framework_version"),
        )


@dataclass
class StateHistory:
    """Historical state snapshot."""
    timestamp: datetime
    state: AgentState
    metadata: StateMetadata
    change_type: str = "update"  # update, create, delete
    
    def to_dict(self) -> Dict[str, Any]:
        return {
            "timestamp": self.timestamp.isoformat(),
            "state": self.state.to_dict(),
            "metadata": self.metadata.to_dict(),
            "change_type": self.change_type,
        }


@dataclass
class UnifiedState:
    """
    Unified state format that all frameworks can read/write.
    """
    session_id: str
    agent_id: str
    state: AgentState
    metadata: StateMetadata
    history: List[StateHistory] = field(default_factory=list)
    
    def to_dict(self) -> Dict[str, Any]:
        """Convert to dictionary for storage."""
        return {
            "session_id": self.session_id,
            "agent_id": self.agent_id,
            "state": self.state.to_dict(),
            "metadata": self.metadata.to_dict(),
            "history": [h.to_dict() for h in self.history],
        }
    
    def to_json(self) -> str:
        """Serialize to JSON string."""
        return json.dumps(self.to_dict(), default=str)
    
    @classmethod
    def from_dict(cls, data: Dict[str, Any]) -> "UnifiedState":
        """Create from dictionary."""
        history = []
        if "history" in data:
            for h in data["history"]:
                history.append(StateHistory(
                    timestamp=datetime.fromisoformat(h["timestamp"]),
                    state=AgentState.from_dict(h["state"]),
                    metadata=StateMetadata.from_dict(h["metadata"]),
                    change_type=h.get("change_type", "update"),
                ))
        
        return cls(
            session_id=data["session_id"],
            agent_id=data["agent_id"],
            state=AgentState.from_dict(data["state"]),
            metadata=StateMetadata.from_dict(data["metadata"]),
            history=history,
        )
    
    @classmethod
    def from_json(cls, json_str: str) -> "UnifiedState":
        """Create from JSON string."""
        return cls.from_dict(json.loads(json_str))
    
    def get_key(self) -> str:
        """Get Redis key for this state."""
        return f"uass:session:{self.session_id}:agent:{self.agent_id}:state"
    
    def get_history_key(self) -> str:
        """Get Redis key for history."""
        return f"uass:session:{self.session_id}:agent:{self.agent_id}:history"
    
    def get_agents_list_key(self) -> str:
        """Get Redis key for agents list."""
        return f"uass:session:{self.session_id}:agents"

