# app/services/memory/episodic/memory.py
"""
Episodic Memory Implementation.

Provides experiential memory storage and retrieval for agent learning.
Combines episode recording and retrieval into a unified interface.
"""

import logging
from datetime import datetime
from typing import Any
from uuid import UUID

from sqlalchemy.ext.asyncio import AsyncSession

from app.services.memory.types import Episode, EpisodeCreate, Outcome, RetrievalResult

from .recorder import EpisodeRecorder
from .retrieval import EpisodeRetriever, RetrievalStrategy

logger = logging.getLogger(__name__)


class EpisodicMemory:
    """
    Episodic Memory Service.

    Provides experiential memory for agent learning:
    - Record task completions with context
    - Store failures with error context
    - Retrieve by semantic similarity
    - Retrieve by recency, outcome, task type
    - Track importance scores
    - Extract lessons learned

    Performance target: <100ms P95 for retrieval
    """

    def __init__(
        self,
        session: AsyncSession,
        embedding_generator: Any | None = None,
    ) -> None:
        """
        Initialize episodic memory.

        Args:
            session: Database session
            embedding_generator: Optional embedding generator for semantic search
        """
        self._session = session
        self._embedding_generator = embedding_generator
        self._recorder = EpisodeRecorder(session, embedding_generator)
        self._retriever = EpisodeRetriever(session, embedding_generator)

    @classmethod
    async def create(
        cls,
        session: AsyncSession,
        embedding_generator: Any | None = None,
    ) -> "EpisodicMemory":
        """
        Factory method to create EpisodicMemory.

        Args:
            session: Database session
            embedding_generator: Optional embedding generator

        Returns:
            Configured EpisodicMemory instance
        """
        return cls(session=session, embedding_generator=embedding_generator)

    # =========================================================================
    # Recording Operations
    # =========================================================================

    async def record_episode(self, episode: EpisodeCreate) -> Episode:
        """
        Record a new episode.

        Args:
            episode: Episode data to record

        Returns:
            The created episode with assigned ID
        """
        return await self._recorder.record(episode)

    async def record_success(
        self,
        project_id: UUID,
        session_id: str,
        task_type: str,
        task_description: str,
        actions: list[dict[str, Any]],
        context_summary: str,
        outcome_details: str = "",
        duration_seconds: float = 0.0,
        tokens_used: int = 0,
        lessons_learned: list[str] | None = None,
        agent_instance_id: UUID | None = None,
        agent_type_id: UUID | None = None,
    ) -> Episode:
        """
        Convenience method to record a successful episode.

        Args:
            project_id: Project ID
            session_id: Session ID
            task_type: Type of task
            task_description: Task description
            actions: Actions taken
            context_summary: Context summary
            outcome_details: Optional outcome details
            duration_seconds: Task duration
            tokens_used: Tokens consumed
            lessons_learned: Optional lessons
            agent_instance_id: Optional agent instance
            agent_type_id: Optional agent type

        Returns:
            The created episode
        """
        episode_data = EpisodeCreate(
            project_id=project_id,
            session_id=session_id,
            task_type=task_type,
            task_description=task_description,
            actions=actions,
            context_summary=context_summary,
            outcome=Outcome.SUCCESS,
            outcome_details=outcome_details,
            duration_seconds=duration_seconds,
            tokens_used=tokens_used,
            lessons_learned=lessons_learned or [],
            agent_instance_id=agent_instance_id,
            agent_type_id=agent_type_id,
        )
        return await self.record_episode(episode_data)

    async def record_failure(
        self,
        project_id: UUID,
        session_id: str,
        task_type: str,
        task_description: str,
        actions: list[dict[str, Any]],
        context_summary: str,
        error_details: str,
        duration_seconds: float = 0.0,
        tokens_used: int = 0,
        lessons_learned: list[str] | None = None,
        agent_instance_id: UUID | None = None,
        agent_type_id: UUID | None = None,
    ) -> Episode:
        """
        Convenience method to record a failed episode.

        Args:
            project_id: Project ID
            session_id: Session ID
            task_type: Type of task
            task_description: Task description
            actions: Actions taken before failure
            context_summary: Context summary
            error_details: Error details
            duration_seconds: Task duration
            tokens_used: Tokens consumed
            lessons_learned: Optional lessons from failure
            agent_instance_id: Optional agent instance
            agent_type_id: Optional agent type

        Returns:
            The created episode
        """
        episode_data = EpisodeCreate(
            project_id=project_id,
            session_id=session_id,
            task_type=task_type,
            task_description=task_description,
            actions=actions,
            context_summary=context_summary,
            outcome=Outcome.FAILURE,
            outcome_details=error_details,
            duration_seconds=duration_seconds,
            tokens_used=tokens_used,
            lessons_learned=lessons_learned or [],
            agent_instance_id=agent_instance_id,
            agent_type_id=agent_type_id,
        )
        return await self.record_episode(episode_data)

    # =========================================================================
    # Retrieval Operations
    # =========================================================================

    async def search_similar(
        self,
        project_id: UUID,
        query: str,
        limit: int = 10,
        agent_instance_id: UUID | None = None,
    ) -> list[Episode]:
        """
        Search for semantically similar episodes.

        Args:
            project_id: Project to search within
            query: Search query
            limit: Maximum results
            agent_instance_id: Optional filter by agent instance

        Returns:
            List of similar episodes
        """
        result = await self._retriever.search_similar(
            project_id, query, limit, agent_instance_id
        )
        return result.items

    async def get_recent(
        self,
        project_id: UUID,
        limit: int = 10,
        since: datetime | None = None,
        agent_instance_id: UUID | None = None,
    ) -> list[Episode]:
        """
        Get recent episodes.

        Args:
            project_id: Project to search within
            limit: Maximum results
            since: Optional time filter
            agent_instance_id: Optional filter by agent instance

        Returns:
            List of recent episodes
        """
        result = await self._retriever.get_recent(
            project_id, limit, since, agent_instance_id
        )
        return result.items

    async def get_by_outcome(
        self,
        project_id: UUID,
        outcome: Outcome,
        limit: int = 10,
        agent_instance_id: UUID | None = None,
    ) -> list[Episode]:
        """
        Get episodes by outcome.

        Args:
            project_id: Project to search within
            outcome: Outcome to filter by
            limit: Maximum results
            agent_instance_id: Optional filter by agent instance

        Returns:
            List of episodes with specified outcome
        """
        result = await self._retriever.get_by_outcome(
            project_id, outcome, limit, agent_instance_id
        )
        return result.items

    async def get_by_task_type(
        self,
        project_id: UUID,
        task_type: str,
        limit: int = 10,
        agent_instance_id: UUID | None = None,
    ) -> list[Episode]:
        """
        Get episodes by task type.

        Args:
            project_id: Project to search within
            task_type: Task type to filter by
            limit: Maximum results
            agent_instance_id: Optional filter by agent instance

        Returns:
            List of episodes with specified task type
        """
        result = await self._retriever.get_by_task_type(
            project_id, task_type, limit, agent_instance_id
        )
        return result.items

    async def get_important(
        self,
        project_id: UUID,
        limit: int = 10,
        min_importance: float = 0.7,
        agent_instance_id: UUID | None = None,
    ) -> list[Episode]:
        """
        Get high-importance episodes.

        Args:
            project_id: Project to search within
            limit: Maximum results
            min_importance: Minimum importance score
            agent_instance_id: Optional filter by agent instance

        Returns:
            List of important episodes
        """
        result = await self._retriever.get_important(
            project_id, limit, min_importance, agent_instance_id
        )
        return result.items

    async def retrieve(
        self,
        project_id: UUID,
        strategy: RetrievalStrategy = RetrievalStrategy.RECENCY,
        limit: int = 10,
        **kwargs: Any,
    ) -> RetrievalResult[Episode]:
        """
        Retrieve episodes with full result metadata.

        Args:
            project_id: Project to search within
            strategy: Retrieval strategy
            limit: Maximum results
            **kwargs: Strategy-specific parameters

        Returns:
            RetrievalResult with episodes and metadata
        """
        return await self._retriever.retrieve(project_id, strategy, limit, **kwargs)

    # =========================================================================
    # Modification Operations
    # =========================================================================

    async def get_by_id(self, episode_id: UUID) -> Episode | None:
        """Get an episode by ID."""
        return await self._recorder.get_by_id(episode_id)

    async def update_importance(
        self,
        episode_id: UUID,
        importance_score: float,
    ) -> Episode | None:
        """
        Update an episode's importance score.

        Args:
            episode_id: Episode to update
            importance_score: New importance score (0.0 to 1.0)

        Returns:
            Updated episode or None if not found
        """
        return await self._recorder.update_importance(episode_id, importance_score)

    async def add_lessons(
        self,
        episode_id: UUID,
        lessons: list[str],
    ) -> Episode | None:
        """
        Add lessons learned to an episode.

        Args:
            episode_id: Episode to update
            lessons: Lessons to add

        Returns:
            Updated episode or None if not found
        """
        return await self._recorder.add_lessons(episode_id, lessons)

    async def delete(self, episode_id: UUID) -> bool:
        """
        Delete an episode.

        Args:
            episode_id: Episode to delete

        Returns:
            True if deleted
        """
        return await self._recorder.delete(episode_id)

    # =========================================================================
    # Summarization
    # =========================================================================

    async def summarize_episodes(
        self,
        episode_ids: list[UUID],
    ) -> str:
        """
        Summarize multiple episodes into a consolidated view.

        Args:
            episode_ids: Episodes to summarize

        Returns:
            Summary text
        """
        if not episode_ids:
            return "No episodes to summarize."

        episodes: list[Episode] = []
        for episode_id in episode_ids:
            episode = await self.get_by_id(episode_id)
            if episode:
                episodes.append(episode)

        if not episodes:
            return "No episodes found."

        # Build summary
        lines = [f"Summary of {len(episodes)} episodes:", ""]

        # Outcome breakdown
        success = sum(1 for e in episodes if e.outcome == Outcome.SUCCESS)
        failure = sum(1 for e in episodes if e.outcome == Outcome.FAILURE)
        partial = sum(1 for e in episodes if e.outcome == Outcome.PARTIAL)
        lines.append(
            f"Outcomes: {success} success, {failure} failure, {partial} partial"
        )

        # Task types
        task_types = {e.task_type for e in episodes}
        lines.append(f"Task types: {', '.join(sorted(task_types))}")

        # Aggregate lessons
        all_lessons: list[str] = []
        for e in episodes:
            all_lessons.extend(e.lessons_learned)

        if all_lessons:
            lines.append("")
            lines.append("Key lessons learned:")
            # Deduplicate lessons
            unique_lessons = list(dict.fromkeys(all_lessons))
            for lesson in unique_lessons[:10]:  # Top 10
                lines.append(f"  - {lesson}")

        # Duration and tokens
        total_duration = sum(e.duration_seconds for e in episodes)
        total_tokens = sum(e.tokens_used for e in episodes)
        lines.append("")
        lines.append(f"Total duration: {total_duration:.1f}s")
        lines.append(f"Total tokens: {total_tokens:,}")

        return "\n".join(lines)

    # =========================================================================
    # Statistics
    # =========================================================================

    async def get_stats(self, project_id: UUID) -> dict[str, Any]:
        """
        Get episode statistics for a project.

        Args:
            project_id: Project to get stats for

        Returns:
            Dictionary with episode statistics
        """
        return await self._recorder.get_stats(project_id)

    async def count(
        self,
        project_id: UUID,
        since: datetime | None = None,
    ) -> int:
        """
        Count episodes for a project.

        Args:
            project_id: Project to count for
            since: Optional time filter

        Returns:
            Number of episodes
        """
        return await self._recorder.count_by_project(project_id, since)
