"""
Memory management for orchestration patterns.

This module provides short-term and long-term memory capabilities for agents
and patterns, including key-value storage, access tracking, and eviction policies.
"""

from typing import Any, Dict, List, Optional
from pydantic import BaseModel, Field
from datetime import datetime


class MemoryEntry(BaseModel):
    """
    Single memory entry with metadata and access tracking.

    This Pydantic model represents an individual piece of stored information
    with timestamp, access count, and custom metadata.

    :param key: str - Unique identifier for this entry.
    :param value: Any - The stored value.
    :param timestamp: datetime - When this entry was created.
    :param access_count: int - Number of times this entry has been accessed.
    :param metadata: Dict[str, Any] - Custom metadata for this entry.
    """

    key: str = Field(..., description="Memory entry key")
    value: Any = Field(..., description="Memory entry value")
    timestamp: datetime = Field(default_factory=datetime.utcnow)
    access_count: int = Field(0, ge=0)
    metadata: Dict[str, Any] = Field(default_factory=dict)

    class Config:
        json_encoders = {datetime: lambda v: v.isoformat()}


class Memory(BaseModel):
    """
    Memory storage with size limits and eviction.

    This Pydantic model provides a bounded key-value store with automatic
    eviction of oldest entries when capacity is reached.

    :param entries: Dict[str, MemoryEntry] - Stored memory entries.
    :param max_size: Optional[int] - Maximum number of entries. None means unlimited.
    """

    entries: Dict[str, MemoryEntry] = Field(default_factory=dict)
    max_size: Optional[int] = Field(None, ge=1, description="Maximum number of entries")

    def store(self, key: str, value: Any, metadata: Optional[Dict[str, Any]] = None) -> None:
        """
        Store a value in memory, evicting oldest if at capacity.

        :param key: str - Entry key.
        :param value: Any - Value to store.
        :param metadata: Optional[Dict[str, Any]] - Custom metadata for this entry.
        """
        if self.max_size and len(self.entries) >= self.max_size and key not in self.entries:
            oldest_key = min(self.entries.keys(), key=lambda k: self.entries[k].timestamp)
            del self.entries[oldest_key]

        self.entries[key] = MemoryEntry(key=key, value=value, metadata=metadata or {})

    def retrieve(self, key: str) -> Optional[Any]:
        """
        Retrieve a value and increment access count.

        :param key: str - Entry key.
        :return: Optional[Any] - Stored value if found, None otherwise.
        """
        entry = self.entries.get(key)
        if entry:
            entry.access_count += 1
            return entry.value
        return None

    def has(self, key: str) -> bool:
        """
        Check if a key exists in memory.

        :param key: str - Entry key.
        :return: bool - True if key exists, False otherwise.
        """
        return key in self.entries

    def remove(self, key: str) -> None:
        """
        Remove an entry from memory.

        :param key: str - Entry key to remove.
        """
        if key in self.entries:
            del self.entries[key]

    def clear(self) -> None:
        """Clear all entries from memory."""
        self.entries.clear()

    def size(self) -> int:
        """
        Get the current number of entries.

        :return: int - Number of stored entries.
        """
        return len(self.entries)

    def get_recent(self, limit: int = 10) -> List[MemoryEntry]:
        """
        Get the most recently created entries.

        :param limit: int - Maximum number of entries to return.
        :return: List[MemoryEntry] - Recent entries sorted by timestamp (newest first).
        """
        sorted_entries = sorted(self.entries.values(), key=lambda e: e.timestamp, reverse=True)
        return sorted_entries[:limit]


class MemoryStore:
    """
    Manager for multiple memory instances.

    Provides centralized management of multiple memory instances,
    enabling separate memory spaces for different agents or contexts.
    """

    def __init__(self, max_size: Optional[int] = None):
        """
        Initialize the memory store.

        :param max_size: Optional[int] - Default max size for created memories.
        """
        self._memories: Dict[str, Memory] = {}
        self._max_size = max_size

    def create_memory(self, memory_id: str, max_size: Optional[int] = None) -> Memory:
        """
        Create and register a new memory instance.

        :param memory_id: str - Unique identifier for the memory.
        :param max_size: Optional[int] - Max size for this memory (overrides default).
        :return: Memory - The newly created memory instance.
        """
        memory = Memory(max_size=max_size or self._max_size)
        self._memories[memory_id] = memory
        return memory

    def get_memory(self, memory_id: str) -> Optional[Memory]:
        """
        Retrieve a memory instance by ID.

        :param memory_id: str - Memory identifier.
        :return: Optional[Memory] - The memory if found, None otherwise.
        """
        return self._memories.get(memory_id)

    def delete_memory(self, memory_id: str) -> None:
        """
        Delete a memory instance.

        :param memory_id: str - Memory identifier to delete.
        """
        if memory_id in self._memories:
            del self._memories[memory_id]

    def list_memories(self) -> list:
        """
        List all registered memory IDs.

        :return: list - List of memory identifiers.
        """
        return list(self._memories.keys())

    def clear_all(self) -> None:
        """Clear all memory instances."""
        self._memories.clear()
