# app/services/memory/episodic/recorder.py
"""
Episode Recording.

Handles the creation and storage of episodic memories
during agent task execution.
"""

import logging
from datetime import UTC, datetime
from typing import Any
from uuid import UUID, uuid4

from sqlalchemy import select, update
from sqlalchemy.ext.asyncio import AsyncSession

from app.models.memory.enums import EpisodeOutcome
from app.models.memory.episode import Episode as EpisodeModel
from app.services.memory.config import get_memory_settings
from app.services.memory.types import Episode, EpisodeCreate, Outcome

logger = logging.getLogger(__name__)


def _outcome_to_db(outcome: Outcome) -> EpisodeOutcome:
    """Convert service Outcome to database EpisodeOutcome."""
    return EpisodeOutcome(outcome.value)


def _db_to_outcome(db_outcome: EpisodeOutcome) -> Outcome:
    """Convert database EpisodeOutcome to service Outcome."""
    return Outcome(db_outcome.value)


def _model_to_episode(model: EpisodeModel) -> Episode:
    """Convert SQLAlchemy model to Episode dataclass."""
    # SQLAlchemy Column types are inferred as Column[T] by mypy, but at runtime
    # they return actual values. We use type: ignore to handle this mismatch.
    return Episode(
        id=model.id,  # type: ignore[arg-type]
        project_id=model.project_id,  # type: ignore[arg-type]
        agent_instance_id=model.agent_instance_id,  # type: ignore[arg-type]
        agent_type_id=model.agent_type_id,  # type: ignore[arg-type]
        session_id=model.session_id,  # type: ignore[arg-type]
        task_type=model.task_type,  # type: ignore[arg-type]
        task_description=model.task_description,  # type: ignore[arg-type]
        actions=model.actions or [],  # type: ignore[arg-type]
        context_summary=model.context_summary,  # type: ignore[arg-type]
        outcome=_db_to_outcome(model.outcome),  # type: ignore[arg-type]
        outcome_details=model.outcome_details or "",  # type: ignore[arg-type]
        duration_seconds=model.duration_seconds,  # type: ignore[arg-type]
        tokens_used=model.tokens_used,  # type: ignore[arg-type]
        lessons_learned=model.lessons_learned or [],  # type: ignore[arg-type]
        importance_score=model.importance_score,  # type: ignore[arg-type]
        embedding=None,  # Don't expose raw embedding
        occurred_at=model.occurred_at,  # type: ignore[arg-type]
        created_at=model.created_at,  # type: ignore[arg-type]
        updated_at=model.updated_at,  # type: ignore[arg-type]
    )


class EpisodeRecorder:
    """
    Records episodes to the database.

    Handles episode creation, importance scoring,
    and lesson extraction.
    """

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

        Args:
            session: Database session
            embedding_generator: Optional embedding generator for semantic indexing
        """
        self._session = session
        self._embedding_generator = embedding_generator
        self._settings = get_memory_settings()

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

        Args:
            episode_data: Episode data to record

        Returns:
            The created episode
        """
        now = datetime.now(UTC)

        # Calculate importance score if not provided
        importance = episode_data.importance_score
        if importance == 0.5:  # Default value, calculate
            importance = self._calculate_importance(episode_data)

        # Create the model
        model = EpisodeModel(
            id=uuid4(),
            project_id=episode_data.project_id,
            agent_instance_id=episode_data.agent_instance_id,
            agent_type_id=episode_data.agent_type_id,
            session_id=episode_data.session_id,
            task_type=episode_data.task_type,
            task_description=episode_data.task_description,
            actions=episode_data.actions,
            context_summary=episode_data.context_summary,
            outcome=_outcome_to_db(episode_data.outcome),
            outcome_details=episode_data.outcome_details,
            duration_seconds=episode_data.duration_seconds,
            tokens_used=episode_data.tokens_used,
            lessons_learned=episode_data.lessons_learned,
            importance_score=importance,
            occurred_at=now,
            created_at=now,
            updated_at=now,
        )

        # Generate embedding if generator available
        if self._embedding_generator is not None:
            try:
                text_for_embedding = self._create_embedding_text(episode_data)
                embedding = await self._embedding_generator.generate(text_for_embedding)
                model.embedding = embedding
            except Exception as e:
                logger.warning(f"Failed to generate embedding: {e}")

        self._session.add(model)
        await self._session.flush()
        await self._session.refresh(model)

        logger.debug(f"Recorded episode {model.id} for task {model.task_type}")
        return _model_to_episode(model)

    def _calculate_importance(self, episode_data: EpisodeCreate) -> float:
        """
        Calculate importance score for an episode.

        Factors:
        - Outcome: Failures are more important to learn from
        - Duration: Longer tasks may be more significant
        - Token usage: Higher usage may indicate complexity
        - Lessons learned: Episodes with lessons are more valuable
        """
        score = 0.5  # Base score

        # Outcome factor
        if episode_data.outcome == Outcome.FAILURE:
            score += 0.2  # Failures are important for learning
        elif episode_data.outcome == Outcome.PARTIAL:
            score += 0.1
        # Success is default, no adjustment

        # Lessons learned factor
        if episode_data.lessons_learned:
            score += min(0.15, len(episode_data.lessons_learned) * 0.05)

        # Duration factor (longer tasks may be more significant)
        if episode_data.duration_seconds > 60:
            score += 0.05
        if episode_data.duration_seconds > 300:
            score += 0.05

        # Token usage factor (complex tasks)
        if episode_data.tokens_used > 1000:
            score += 0.05

        # Clamp to valid range
        return min(1.0, max(0.0, score))

    def _create_embedding_text(self, episode_data: EpisodeCreate) -> str:
        """Create text representation for embedding generation."""
        parts = [
            f"Task: {episode_data.task_type}",
            f"Description: {episode_data.task_description}",
            f"Context: {episode_data.context_summary}",
            f"Outcome: {episode_data.outcome.value}",
        ]

        if episode_data.outcome_details:
            parts.append(f"Details: {episode_data.outcome_details}")

        if episode_data.lessons_learned:
            parts.append(f"Lessons: {', '.join(episode_data.lessons_learned)}")

        return "\n".join(parts)

    async def get_by_id(self, episode_id: UUID) -> Episode | None:
        """Get an episode by ID."""
        query = select(EpisodeModel).where(EpisodeModel.id == episode_id)
        result = await self._session.execute(query)
        model = result.scalar_one_or_none()

        if model is None:
            return None

        return _model_to_episode(model)

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

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

        Returns:
            Updated episode or None if not found
        """
        # Validate score
        importance_score = min(1.0, max(0.0, importance_score))

        stmt = (
            update(EpisodeModel)
            .where(EpisodeModel.id == episode_id)
            .values(
                importance_score=importance_score,
                updated_at=datetime.now(UTC),
            )
            .returning(EpisodeModel)
        )

        result = await self._session.execute(stmt)
        model = result.scalar_one_or_none()

        if model is None:
            return None

        await self._session.flush()
        return _model_to_episode(model)

    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: New lessons to add

        Returns:
            Updated episode or None if not found
        """
        # Get current episode
        query = select(EpisodeModel).where(EpisodeModel.id == episode_id)
        result = await self._session.execute(query)
        model = result.scalar_one_or_none()

        if model is None:
            return None

        # Append lessons
        current_lessons: list[str] = model.lessons_learned or []  # type: ignore[assignment]
        updated_lessons = current_lessons + lessons

        stmt = (
            update(EpisodeModel)
            .where(EpisodeModel.id == episode_id)
            .values(
                lessons_learned=updated_lessons,
                updated_at=datetime.now(UTC),
            )
            .returning(EpisodeModel)
        )

        result = await self._session.execute(stmt)
        model = result.scalar_one_or_none()
        await self._session.flush()

        return _model_to_episode(model) if model else None

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

        Args:
            episode_id: Episode to delete

        Returns:
            True if deleted
        """
        query = select(EpisodeModel).where(EpisodeModel.id == episode_id)
        result = await self._session.execute(query)
        model = result.scalar_one_or_none()

        if model is None:
            return False

        await self._session.delete(model)
        await self._session.flush()
        return True

    async def count_by_project(
        self,
        project_id: UUID,
        since: datetime | None = None,
    ) -> int:
        """Count episodes for a project."""
        query = select(EpisodeModel).where(EpisodeModel.project_id == project_id)
        if since is not None:
            query = query.where(EpisodeModel.occurred_at >= since)

        result = await self._session.execute(query)
        return len(list(result.scalars().all()))

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

        Returns:
            Dictionary with episode statistics
        """
        query = select(EpisodeModel).where(EpisodeModel.project_id == project_id)
        result = await self._session.execute(query)
        episodes = list(result.scalars().all())

        if not episodes:
            return {
                "total_count": 0,
                "success_count": 0,
                "failure_count": 0,
                "partial_count": 0,
                "avg_importance": 0.0,
                "avg_duration": 0.0,
                "total_tokens": 0,
            }

        success_count = sum(1 for e in episodes if e.outcome == EpisodeOutcome.SUCCESS)
        failure_count = sum(1 for e in episodes if e.outcome == EpisodeOutcome.FAILURE)
        partial_count = sum(1 for e in episodes if e.outcome == EpisodeOutcome.PARTIAL)

        avg_importance = sum(e.importance_score for e in episodes) / len(episodes)
        avg_duration = sum(e.duration_seconds for e in episodes) / len(episodes)
        total_tokens = sum(e.tokens_used for e in episodes)

        return {
            "total_count": len(episodes),
            "success_count": success_count,
            "failure_count": failure_count,
            "partial_count": partial_count,
            "avg_importance": avg_importance,
            "avg_duration": avg_duration,
            "total_tokens": total_tokens,
        }
