"""Agent memory management — short, medium, and long-term memory layers."""

from __future__ import annotations

import logging
from datetime import UTC, datetime
from enum import StrEnum

from pydantic import BaseModel, Field

logger = logging.getLogger(__name__)


class MemoryType(StrEnum):
    SHORT_TERM = "short_term"  # Current conversation context
    MEDIUM_TERM = "medium_term"  # Session-level (persists across conversations in a session)
    LONG_TERM = "long_term"  # Permanent knowledge (stored in VKB)


class MemoryEntry(BaseModel):
    """A single memory entry."""

    id: str
    agent_id: str
    memory_type: MemoryType
    content: str
    tags: list[str] = Field(default_factory=list)
    source: str = ""  # where this memory came from
    confidence: float = 1.0  # how confident the agent is in this memory
    timestamp: datetime = Field(default_factory=lambda: datetime.now(UTC))
    metadata: dict = Field(default_factory=dict)


class AgentMemory:
    """
    Three-tier memory system for an agent.

    - Short-term: in-memory list, auto-pruned (conversation context)
    - Medium-term: in-memory dict, persists for session duration
    - Long-term: backed by KnowledgeStore (VKB)
    """

    def __init__(
        self,
        agent_id: str,
        knowledge_store: KnowledgeStore | None = None,  # noqa: F821
        short_term_limit: int = 50,
        medium_term_limit: int = 200,
    ) -> None:
        self.agent_id = agent_id
        self.knowledge_store = knowledge_store
        self.short_term_limit = short_term_limit
        self.medium_term_limit = medium_term_limit

        self._short_term: list[MemoryEntry] = []
        self._medium_term: dict[str, MemoryEntry] = {}

    # ── Short-Term ───────────────────────────────────

    def remember_short(
        self, content: str, tags: list[str] | None = None, **kwargs: object
    ) -> MemoryEntry:
        """Add an entry to short-term memory."""
        entry = MemoryEntry(
            id=f"st-{len(self._short_term)}",
            agent_id=self.agent_id,
            memory_type=MemoryType.SHORT_TERM,
            content=content,
            tags=tags or [],
            **kwargs,
        )
        self._short_term.append(entry)
        if len(self._short_term) > self.short_term_limit:
            self._short_term = self._short_term[-self.short_term_limit :]
        return entry

    def get_short_term(self, limit: int = 10) -> list[MemoryEntry]:
        """Retrieve recent short-term memories."""
        return self._short_term[-limit:]

    # ── Medium-Term ──────────────────────────────────

    def remember_medium(
        self, key: str, content: str, tags: list[str] | None = None, **kwargs: object
    ) -> MemoryEntry:
        """Store a keyed medium-term memory (session-level)."""
        entry = MemoryEntry(
            id=f"mt-{key}",
            agent_id=self.agent_id,
            memory_type=MemoryType.MEDIUM_TERM,
            content=content,
            tags=tags or [],
            **kwargs,
        )
        self._medium_term[key] = entry
        # Evict oldest if over limit
        if len(self._medium_term) > self.medium_term_limit:
            oldest_key = next(iter(self._medium_term))
            del self._medium_term[oldest_key]
        return entry

    def recall_medium(self, key: str) -> MemoryEntry | None:
        """Retrieve a specific medium-term memory by key."""
        return self._medium_term.get(key)

    def search_medium(self, tags: list[str]) -> list[MemoryEntry]:
        """Search medium-term memories by tags."""
        tag_set = set(tags)
        return [m for m in self._medium_term.values() if tag_set.intersection(m.tags)]

    # ── Long-Term (VKB) ─────────────────────────────

    async def remember_long(
        self, content: str, tags: list[str] | None = None, **kwargs: object
    ) -> str | None:
        """Persist a memory to the long-term Knowledge Store."""
        if self.knowledge_store is None:
            logger.warning("[%s] No knowledge store — long-term memory not saved", self.agent_id)
            return None
        entry_id = await self.knowledge_store.store_entry(
            content=content,
            author=self.agent_id,
            tags=tags or [],
            metadata=kwargs,
        )
        return entry_id

    async def recall_long(self, query: str, top_k: int = 5) -> list[dict]:
        """Search long-term memory via semantic similarity."""
        if self.knowledge_store is None:
            return []
        return await self.knowledge_store.search(query=query, top_k=top_k)

    # ── Context Assembly ─────────────────────────────

    def get_context_window(self, max_entries: int = 20) -> str:
        """Build a text context window from recent memories across tiers."""
        parts: list[str] = []

        # Medium-term summaries
        medium_entries = list(self._medium_term.values())[-5:]
        if medium_entries:
            parts.append("## Session Context")
            for entry in medium_entries:
                parts.append(f"- [{', '.join(entry.tags)}] {entry.content[:200]}")

        # Short-term recent
        short_entries = self._short_term[-max_entries:]
        if short_entries:
            parts.append("\n## Recent Context")
            for entry in short_entries:
                parts.append(f"- {entry.content[:200]}")

        return "\n".join(parts)

    def clear_short_term(self) -> None:
        """Clear short-term memory (e.g., start of a new conversation)."""
        self._short_term.clear()

    def clear_session(self) -> None:
        """Clear both short-term and medium-term memory."""
        self._short_term.clear()
        self._medium_term.clear()
