"""
Memory System Types

Core type definitions and interfaces for the Agent Memory System.
"""

from abc import ABC, abstractmethod
from dataclasses import dataclass, field
from datetime import UTC, datetime
from enum import Enum
from typing import Any
from uuid import UUID


def _utcnow() -> datetime:
    """Get current UTC time as timezone-aware datetime."""
    return datetime.now(UTC)


class MemoryType(str, Enum):
    """Types of memory in the agent memory system."""

    WORKING = "working"
    EPISODIC = "episodic"
    SEMANTIC = "semantic"
    PROCEDURAL = "procedural"


class ScopeLevel(str, Enum):
    """Hierarchical scoping levels for memory."""

    GLOBAL = "global"
    PROJECT = "project"
    AGENT_TYPE = "agent_type"
    AGENT_INSTANCE = "agent_instance"
    SESSION = "session"


class Outcome(str, Enum):
    """Outcome of a task or episode."""

    SUCCESS = "success"
    FAILURE = "failure"
    PARTIAL = "partial"
    ABANDONED = "abandoned"


class ConsolidationStatus(str, Enum):
    """Status of a memory consolidation job."""

    PENDING = "pending"
    RUNNING = "running"
    COMPLETED = "completed"
    FAILED = "failed"


class ConsolidationType(str, Enum):
    """Types of memory consolidation."""

    WORKING_TO_EPISODIC = "working_to_episodic"
    EPISODIC_TO_SEMANTIC = "episodic_to_semantic"
    EPISODIC_TO_PROCEDURAL = "episodic_to_procedural"
    PRUNING = "pruning"


@dataclass
class ScopeContext:
    """Represents a memory scope with its hierarchy."""

    scope_type: ScopeLevel
    scope_id: str
    parent: "ScopeContext | None" = None

    def get_hierarchy(self) -> list["ScopeContext"]:
        """Get the full scope hierarchy from root to this scope."""
        hierarchy: list[ScopeContext] = []
        current: ScopeContext | None = self
        while current is not None:
            hierarchy.insert(0, current)
            current = current.parent
        return hierarchy

    def to_key_prefix(self) -> str:
        """Convert scope to a key prefix for storage."""
        return f"{self.scope_type.value}:{self.scope_id}"


@dataclass
class MemoryItem:
    """Base class for all memory items."""

    id: UUID
    memory_type: MemoryType
    scope_type: ScopeLevel
    scope_id: str
    created_at: datetime
    updated_at: datetime
    metadata: dict[str, Any] = field(default_factory=dict)

    def get_age_seconds(self) -> float:
        """Get the age of this memory item in seconds."""
        return (_utcnow() - self.created_at).total_seconds()


@dataclass
class WorkingMemoryItem:
    """A key-value item in working memory."""

    id: UUID
    scope_type: ScopeLevel
    scope_id: str
    key: str
    value: Any
    expires_at: datetime | None = None
    created_at: datetime = field(default_factory=_utcnow)
    updated_at: datetime = field(default_factory=_utcnow)

    def is_expired(self) -> bool:
        """Check if this item has expired."""
        if self.expires_at is None:
            return False
        return _utcnow() > self.expires_at


@dataclass
class TaskState:
    """Current state of a task in working memory."""

    task_id: str
    task_type: str
    description: str
    status: str = "in_progress"
    current_step: int = 0
    total_steps: int = 0
    progress_percent: float = 0.0
    context: dict[str, Any] = field(default_factory=dict)
    started_at: datetime = field(default_factory=_utcnow)
    updated_at: datetime = field(default_factory=_utcnow)


@dataclass
class Episode:
    """An episodic memory - a recorded experience."""

    id: UUID
    project_id: UUID
    agent_instance_id: UUID | None
    agent_type_id: UUID | None
    session_id: str
    task_type: str
    task_description: str
    actions: list[dict[str, Any]]
    context_summary: str
    outcome: Outcome
    outcome_details: str
    duration_seconds: float
    tokens_used: int
    lessons_learned: list[str]
    importance_score: float
    embedding: list[float] | None
    occurred_at: datetime
    created_at: datetime
    updated_at: datetime


@dataclass
class EpisodeCreate:
    """Data required to create a new episode."""

    project_id: UUID
    session_id: str
    task_type: str
    task_description: str
    actions: list[dict[str, Any]]
    context_summary: str
    outcome: Outcome
    outcome_details: str
    duration_seconds: float
    tokens_used: int
    lessons_learned: list[str] = field(default_factory=list)
    importance_score: float = 0.5
    agent_instance_id: UUID | None = None
    agent_type_id: UUID | None = None


@dataclass
class Fact:
    """A semantic memory fact - a piece of knowledge."""

    id: UUID
    project_id: UUID | None  # None for global facts
    subject: str
    predicate: str
    object: str
    confidence: float
    source_episode_ids: list[UUID]
    first_learned: datetime
    last_reinforced: datetime
    reinforcement_count: int
    embedding: list[float] | None
    created_at: datetime
    updated_at: datetime


@dataclass
class FactCreate:
    """Data required to create a new fact."""

    subject: str
    predicate: str
    object: str
    confidence: float = 0.8
    project_id: UUID | None = None
    source_episode_ids: list[UUID] = field(default_factory=list)


@dataclass
class Procedure:
    """A procedural memory - a learned skill or procedure."""

    id: UUID
    project_id: UUID | None
    agent_type_id: UUID | None
    name: str
    trigger_pattern: str
    steps: list[dict[str, Any]]
    success_count: int
    failure_count: int
    last_used: datetime | None
    embedding: list[float] | None
    created_at: datetime
    updated_at: datetime

    @property
    def success_rate(self) -> float:
        """Calculate the success rate of this procedure."""
        total = self.success_count + self.failure_count
        if total == 0:
            return 0.0
        return self.success_count / total


@dataclass
class ProcedureCreate:
    """Data required to create a new procedure."""

    name: str
    trigger_pattern: str
    steps: list[dict[str, Any]]
    project_id: UUID | None = None
    agent_type_id: UUID | None = None


@dataclass
class Step:
    """A single step in a procedure."""

    order: int
    action: str
    parameters: dict[str, Any] = field(default_factory=dict)
    expected_outcome: str = ""
    fallback_action: str | None = None


class MemoryStore[T: MemoryItem](ABC):
    """Abstract base class for memory storage backends."""

    @abstractmethod
    async def store(self, item: T) -> T:
        """Store a memory item."""
        ...

    @abstractmethod
    async def get(self, item_id: UUID) -> T | None:
        """Get a memory item by ID."""
        ...

    @abstractmethod
    async def delete(self, item_id: UUID) -> bool:
        """Delete a memory item."""
        ...

    @abstractmethod
    async def list(
        self,
        scope_type: ScopeLevel | None = None,
        scope_id: str | None = None,
        limit: int = 100,
        offset: int = 0,
    ) -> list[T]:
        """List memory items with optional scope filtering."""
        ...

    @abstractmethod
    async def count(
        self,
        scope_type: ScopeLevel | None = None,
        scope_id: str | None = None,
    ) -> int:
        """Count memory items with optional scope filtering."""
        ...


@dataclass
class RetrievalResult[T]:
    """Result of a memory retrieval operation."""

    items: list[T]
    total_count: int
    query: str
    retrieval_type: str
    latency_ms: float
    metadata: dict[str, Any] = field(default_factory=dict)


@dataclass
class MemoryStats:
    """Statistics about memory usage."""

    memory_type: MemoryType
    scope_type: ScopeLevel | None
    scope_id: str | None
    item_count: int
    total_size_bytes: int
    oldest_item_age_seconds: float
    newest_item_age_seconds: float
    avg_item_size_bytes: float
    metadata: dict[str, Any] = field(default_factory=dict)
