# app/services/memory/integration/context_source.py
"""
Memory Context Source.

Provides agent memory as a context source for the Context Engine.
Retrieves relevant memories based on query and converts them to MemoryContext objects.
"""

import logging
from dataclasses import dataclass
from datetime import UTC, datetime, timedelta
from typing import Any
from uuid import UUID

from sqlalchemy.ext.asyncio import AsyncSession

from app.services.context.types.memory import MemoryContext
from app.services.memory.episodic import EpisodicMemory
from app.services.memory.procedural import ProceduralMemory
from app.services.memory.semantic import SemanticMemory
from app.services.memory.working import WorkingMemory

logger = logging.getLogger(__name__)


@dataclass
class MemoryFetchConfig:
    """Configuration for memory fetching."""

    # Limits per memory type
    working_limit: int = 10
    episodic_limit: int = 10
    semantic_limit: int = 15
    procedural_limit: int = 5

    # Time ranges
    episodic_days_back: int = 30
    min_relevance: float = 0.3

    # Which memory types to include
    include_working: bool = True
    include_episodic: bool = True
    include_semantic: bool = True
    include_procedural: bool = True


@dataclass
class MemoryFetchResult:
    """Result of memory fetch operation."""

    contexts: list[MemoryContext]
    by_type: dict[str, int]
    fetch_time_ms: float
    query: str


class MemoryContextSource:
    """
    Source for memory context in the Context Engine.

    This service retrieves relevant memories based on a query and
    converts them to MemoryContext objects for context assembly.
    It coordinates between all memory types (working, episodic,
    semantic, procedural) to provide a comprehensive memory context.
    """

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

        Args:
            session: Database session
            embedding_generator: Optional embedding generator for semantic search
        """
        self._session = session
        self._embedding_generator = embedding_generator

        # Lazy-initialized memory services
        self._episodic: EpisodicMemory | None = None
        self._semantic: SemanticMemory | None = None
        self._procedural: ProceduralMemory | None = None

    async def _get_episodic(self) -> EpisodicMemory:
        """Get or create episodic memory service."""
        if self._episodic is None:
            self._episodic = await EpisodicMemory.create(
                self._session,
                self._embedding_generator,
            )
        return self._episodic

    async def _get_semantic(self) -> SemanticMemory:
        """Get or create semantic memory service."""
        if self._semantic is None:
            self._semantic = await SemanticMemory.create(
                self._session,
                self._embedding_generator,
            )
        return self._semantic

    async def _get_procedural(self) -> ProceduralMemory:
        """Get or create procedural memory service."""
        if self._procedural is None:
            self._procedural = await ProceduralMemory.create(
                self._session,
                self._embedding_generator,
            )
        return self._procedural

    async def fetch_context(
        self,
        query: str,
        project_id: UUID,
        agent_instance_id: UUID | None = None,
        agent_type_id: UUID | None = None,
        session_id: str | None = None,
        config: MemoryFetchConfig | None = None,
    ) -> MemoryFetchResult:
        """
        Fetch relevant memories as context.

        This is the main entry point for the Context Engine integration.
        It searches across all memory types and returns relevant memories
        as MemoryContext objects.

        Args:
            query: Search query for finding relevant memories
            project_id: Project scope
            agent_instance_id: Optional agent instance scope
            agent_type_id: Optional agent type scope (for procedural)
            session_id: Optional session ID (for working memory)
            config: Optional fetch configuration

        Returns:
            MemoryFetchResult with contexts and metadata
        """
        config = config or MemoryFetchConfig()
        start_time = datetime.now(UTC)

        contexts: list[MemoryContext] = []
        by_type: dict[str, int] = {
            "working": 0,
            "episodic": 0,
            "semantic": 0,
            "procedural": 0,
        }

        # Fetch from working memory (session-scoped)
        if config.include_working and session_id:
            try:
                working_contexts = await self._fetch_working(
                    query=query,
                    session_id=session_id,
                    project_id=project_id,
                    agent_instance_id=agent_instance_id,
                    limit=config.working_limit,
                )
                contexts.extend(working_contexts)
                by_type["working"] = len(working_contexts)
            except Exception as e:
                logger.warning(f"Failed to fetch working memory: {e}")

        # Fetch from episodic memory
        if config.include_episodic:
            try:
                episodic_contexts = await self._fetch_episodic(
                    query=query,
                    project_id=project_id,
                    agent_instance_id=agent_instance_id,
                    limit=config.episodic_limit,
                    days_back=config.episodic_days_back,
                )
                contexts.extend(episodic_contexts)
                by_type["episodic"] = len(episodic_contexts)
            except Exception as e:
                logger.warning(f"Failed to fetch episodic memory: {e}")

        # Fetch from semantic memory
        if config.include_semantic:
            try:
                semantic_contexts = await self._fetch_semantic(
                    query=query,
                    project_id=project_id,
                    limit=config.semantic_limit,
                    min_relevance=config.min_relevance,
                )
                contexts.extend(semantic_contexts)
                by_type["semantic"] = len(semantic_contexts)
            except Exception as e:
                logger.warning(f"Failed to fetch semantic memory: {e}")

        # Fetch from procedural memory
        if config.include_procedural:
            try:
                procedural_contexts = await self._fetch_procedural(
                    query=query,
                    project_id=project_id,
                    agent_type_id=agent_type_id,
                    limit=config.procedural_limit,
                )
                contexts.extend(procedural_contexts)
                by_type["procedural"] = len(procedural_contexts)
            except Exception as e:
                logger.warning(f"Failed to fetch procedural memory: {e}")

        # Sort by relevance
        contexts.sort(key=lambda c: c.relevance_score, reverse=True)

        fetch_time = (datetime.now(UTC) - start_time).total_seconds() * 1000

        logger.debug(
            f"Fetched {len(contexts)} memory contexts for query '{query[:50]}...' "
            f"in {fetch_time:.1f}ms"
        )

        return MemoryFetchResult(
            contexts=contexts,
            by_type=by_type,
            fetch_time_ms=fetch_time,
            query=query,
        )

    async def _fetch_working(
        self,
        query: str,
        session_id: str,
        project_id: UUID,
        agent_instance_id: UUID | None,
        limit: int,
    ) -> list[MemoryContext]:
        """Fetch from working memory."""
        working = await WorkingMemory.for_session(
            session_id=session_id,
            project_id=str(project_id),
            agent_instance_id=str(agent_instance_id) if agent_instance_id else None,
        )

        contexts: list[MemoryContext] = []
        all_keys = await working.list_keys()

        # Filter keys by query (simple substring match)
        query_lower = query.lower()
        matched_keys = [k for k in all_keys if query_lower in k.lower()]

        # If no query match, include all keys (working memory is always relevant)
        if not matched_keys and query:
            matched_keys = all_keys

        for key in matched_keys[:limit]:
            value = await working.get(key)
            if value is not None:
                contexts.append(
                    MemoryContext.from_working_memory(
                        key=key,
                        value=value,
                        source=f"working:{session_id}",
                        query=query,
                    )
                )

        return contexts

    async def _fetch_episodic(
        self,
        query: str,
        project_id: UUID,
        agent_instance_id: UUID | None,
        limit: int,
        days_back: int,
    ) -> list[MemoryContext]:
        """Fetch from episodic memory."""
        episodic = await self._get_episodic()

        # Search for similar episodes
        episodes = await episodic.search_similar(
            project_id=project_id,
            query=query,
            limit=limit,
            agent_instance_id=agent_instance_id,
        )

        # Also get recent episodes if we didn't find enough
        if len(episodes) < limit // 2:
            since = datetime.now(UTC) - timedelta(days=days_back)
            recent = await episodic.get_recent(
                project_id=project_id,
                limit=limit,
                since=since,
            )
            # Deduplicate by ID
            existing_ids = {e.id for e in episodes}
            for ep in recent:
                if ep.id not in existing_ids:
                    episodes.append(ep)
                    if len(episodes) >= limit:
                        break

        return [
            MemoryContext.from_episodic_memory(ep, query=query)
            for ep in episodes[:limit]
        ]

    async def _fetch_semantic(
        self,
        query: str,
        project_id: UUID,
        limit: int,
        min_relevance: float,
    ) -> list[MemoryContext]:
        """Fetch from semantic memory."""
        semantic = await self._get_semantic()

        facts = await semantic.search_facts(
            query=query,
            project_id=project_id,
            limit=limit,
            min_confidence=min_relevance,
        )

        return [MemoryContext.from_semantic_memory(fact, query=query) for fact in facts]

    async def _fetch_procedural(
        self,
        query: str,
        project_id: UUID,
        agent_type_id: UUID | None,
        limit: int,
    ) -> list[MemoryContext]:
        """Fetch from procedural memory."""
        procedural = await self._get_procedural()

        procedures = await procedural.find_matching(
            context=query,
            project_id=project_id,
            agent_type_id=agent_type_id,
            limit=limit,
        )

        return [
            MemoryContext.from_procedural_memory(proc, query=query)
            for proc in procedures
        ]

    async def fetch_all_working(
        self,
        session_id: str,
        project_id: UUID,
        agent_instance_id: UUID | None = None,
    ) -> list[MemoryContext]:
        """
        Fetch all working memory for a session.

        Useful for including entire session state in context.

        Args:
            session_id: Session ID
            project_id: Project scope
            agent_instance_id: Optional agent instance scope

        Returns:
            List of MemoryContext for all working memory items
        """
        working = await WorkingMemory.for_session(
            session_id=session_id,
            project_id=str(project_id),
            agent_instance_id=str(agent_instance_id) if agent_instance_id else None,
        )

        contexts: list[MemoryContext] = []
        all_keys = await working.list_keys()

        for key in all_keys:
            value = await working.get(key)
            if value is not None:
                contexts.append(
                    MemoryContext.from_working_memory(
                        key=key,
                        value=value,
                        source=f"working:{session_id}",
                    )
                )

        return contexts


# Factory function
async def get_memory_context_source(
    session: AsyncSession,
    embedding_generator: Any | None = None,
) -> MemoryContextSource:
    """Create a memory context source instance."""
    return MemoryContextSource(
        session=session,
        embedding_generator=embedding_generator,
    )
