"""
H-MEM Level 0: Episode dataclass for agent experience storage.

Episodes are discrete units of agent experience that capture
context, action, and outcome for later retrieval and consolidation.
"""

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


def _utc_now_iso() -> str:
    """Get current UTC time as ISO format string."""
    return datetime.now(timezone.utc).isoformat()


class EpisodeType(Enum):
    """Types of agent episodes for categorization."""

    TASK_EXECUTION = "task_execution"
    PROBLEM_SOLVING = "problem_solving"
    ERROR_RECOVERY = "error_recovery"
    LEARNING = "learning"
    COLLABORATION = "collaboration"


@dataclass
class Episode:
    """
    A discrete agent experience episode (H-MEM Level 0).

    Episodes capture the full context of an agent's action including:
    - What led to the action (context)
    - What was done (action)
    - What happened (outcome)

    This forms the foundation for memory consolidation and semantic search.
    """

    # Required identification fields
    episode_id: str
    agent_id: str
    session_id: str
    episode_type: EpisodeType

    # Core experience fields
    context: str = ""
    action: str = ""
    outcome: str = ""
    success: bool = False

    # Metadata fields
    timestamp: str = field(default_factory=_utc_now_iso)
    tokens_used: int = 0
    duration_ms: int = 0

    # Optional fields for advanced features
    embedding: Optional[List[float]] = None
    tags: List[str] = field(default_factory=list)
    trace_id: Optional[str] = None  # For consolidation linkage

    def to_dict(self) -> Dict[str, Any]:
        """
        Convert episode to dictionary for serialization.

        Returns:
            Dictionary representation suitable for JSON storage.
        """
        return {
            "episode_id": self.episode_id,
            "agent_id": self.agent_id,
            "session_id": self.session_id,
            "episode_type": self.episode_type.value,
            "context": self.context,
            "action": self.action,
            "outcome": self.outcome,
            "success": self.success,
            "timestamp": self.timestamp,
            "tokens_used": self.tokens_used,
            "duration_ms": self.duration_ms,
            "embedding": self.embedding,
            "tags": self.tags,
            "trace_id": self.trace_id,
        }

    @classmethod
    def from_dict(cls, data: Dict[str, Any]) -> "Episode":
        """
        Create Episode from dictionary.

        Args:
            data: Dictionary with episode fields.

        Returns:
            Episode instance.
        """
        # Handle episode_type conversion from string
        episode_type = data.get("episode_type", EpisodeType.TASK_EXECUTION)
        if isinstance(episode_type, str):
            episode_type = EpisodeType(episode_type)

        return cls(
            episode_id=data.get("episode_id", ""),
            agent_id=data.get("agent_id", ""),
            session_id=data.get("session_id", ""),
            episode_type=episode_type,
            context=data.get("context", ""),
            action=data.get("action", ""),
            outcome=data.get("outcome", ""),
            success=data.get("success", False),
            timestamp=data.get("timestamp", _utc_now_iso()),
            tokens_used=data.get("tokens_used", 0),
            duration_ms=data.get("duration_ms", 0),
            embedding=data.get("embedding"),
            tags=data.get("tags", []),
            trace_id=data.get("trace_id"),
        )

    def __post_init__(self):
        """Validate and normalize fields after initialization."""
        # Ensure tags is always a list
        if self.tags is None:
            self.tags = []

        # Ensure timestamp has a value
        if not self.timestamp:
            self.timestamp = _utc_now_iso()

    @property
    def is_consolidated(self) -> bool:
        """Check if episode has been consolidated into higher-level memory."""
        return self.trace_id is not None

    @property
    def has_embedding(self) -> bool:
        """Check if episode has embedding for semantic search."""
        return self.embedding is not None and len(self.embedding) > 0

    def add_tag(self, tag: str) -> None:
        """Add a tag to the episode if not already present."""
        if tag not in self.tags:
            self.tags.append(tag)

    def remove_tag(self, tag: str) -> bool:
        """Remove a tag from the episode. Returns True if tag was present."""
        if tag in self.tags:
            self.tags.remove(tag)
            return True
        return False

    def set_embedding(self, embedding: List[float]) -> None:
        """Set the embedding vector for semantic search."""
        self.embedding = embedding

    def mark_consolidated(self, trace_id: str) -> None:
        """Mark episode as consolidated with given trace ID."""
        self.trace_id = trace_id
