"""
Context Assembler for GSD-RLM Memory System.

Combines MemoryRouter with actual data retrieval to produce
unified context from multiple memory stores for LLM consumption.
"""

from dataclasses import dataclass, field
from typing import Any, Dict, List, Optional, TYPE_CHECKING

if TYPE_CHECKING:
    from gsd_rlm.session.memory import FileSessionMemory
    from gsd_rlm.memory.hmem.store import EpisodeStore
    from gsd_rlm.memory.bridge.store import MemoryBridge
    from gsd_rlm.memory.bridge.router import MemoryRouter, MemoryStoreType
    from gsd_rlm.memory.hmem.retrieval import HMEMRetrieval


@dataclass
class AssembledContext:
    """
    Assembled context from multiple memory stores.

    Contains formatted data from various memory sources ready for
    LLM consumption.
    """

    session: Dict[str, Any] = field(default_factory=dict)
    recent_episodes: List[Dict[str, Any]] = field(default_factory=list)
    project_facts: Dict[str, Any] = field(default_factory=dict)
    relevant_traces: List[Dict[str, Any]] = field(default_factory=list)
    bridge_facts: Dict[str, Any] = field(default_factory=dict)
    routes_used: List[str] = field(default_factory=list)

    def to_dict(self) -> Dict[str, Any]:
        """Convert to dictionary for serialization."""
        return {
            "session": self.session,
            "recent_episodes": self.recent_episodes,
            "project_facts": self.project_facts,
            "relevant_traces": self.relevant_traces,
            "bridge_facts": self.bridge_facts,
            "routes_used": self.routes_used,
        }

    def to_llm_context(self) -> str:
        """
        Format context as a string suitable for LLM prompting.

        Returns:
            Formatted context string.
        """
        parts: List[str] = []

        if self.session:
            parts.append("## Current Session")
            if "messages" in self.session:
                parts.append("Recent messages:")
                for msg in self.session["messages"][-5:]:
                    parts.append(f"  - {msg}")
            if "recent_tasks" in self.session:
                parts.append("Recent tasks:")
                for task in self.session["recent_tasks"][-3:]:
                    parts.append(f"  - {task}")

        if self.recent_episodes:
            parts.append("\n## Recent Episodes")
            for ep in self.recent_episodes[:5]:
                action = ep.get("action", "Unknown action")
                outcome = ep.get("outcome", "")
                success = "✓" if ep.get("success") else "✗"
                parts.append(f"  [{success}] {action}")
                if outcome:
                    parts.append(f"      {outcome[:100]}")

        if self.project_facts:
            parts.append("\n## Project Facts")
            for key, value in list(self.project_facts.items())[:10]:
                parts.append(f"  - {key}: {value}")

        if self.relevant_traces:
            parts.append("\n## Relevant Work Sessions")
            for trace in self.relevant_traces[:3]:
                theme = trace.get("theme", "Unknown theme")
                summary = trace.get("summary", "")
                parts.append(f"  - {theme}")
                if summary:
                    parts.append(f"      {summary[:150]}")

        if self.bridge_facts:
            parts.append("\n## Bridge Facts")
            for level, facts in self.bridge_facts.items():
                if facts:
                    parts.append(f"  [{level}]")
                    for key, value in list(facts.items())[:5]:
                        parts.append(f"    - {key}: {value}")

        return "\n".join(parts) if parts else "No relevant context found."


class ContextAssembler:
    """
    Assembles unified context from multiple memory stores.

    Uses MemoryRouter to determine which stores to query, then
    retrieves and formats data from each relevant store.

    Attributes:
        session_memory: FileSessionMemory for current session.
        episode_store: EpisodeStore for H-MEM episodes.
        memory_bridge: MemoryBridge for persistent facts.
        memory_router: MemoryRouter for routing decisions.
        hmem_retrieval: HMEMRetrieval for semantic search.

    Example:
        >>> assembler = ContextAssembler(
        ...     session_memory=session,
        ...     episode_store=episodes,
        ...     memory_bridge=bridge,
        ...     memory_router=router,
        ...     hmem_retrieval=retrieval,
        ... )
        >>> context = await assembler.get_context_for_task("sess-001", "fix bug")
    """

    def __init__(
        self,
        session_memory: Optional["FileSessionMemory"] = None,
        episode_store: Optional["EpisodeStore"] = None,
        memory_bridge: Optional["MemoryBridge"] = None,
        memory_router: Optional["MemoryRouter"] = None,
        hmem_retrieval: Optional["HMEMRetrieval"] = None,
    ):
        """
        Initialize ContextAssembler.

        Args:
            session_memory: Optional session memory for current context.
            episode_store: Optional episode store for H-MEM retrieval.
            memory_bridge: Optional memory bridge for persistent facts.
            memory_router: Optional memory router for routing decisions.
            hmem_retrieval: Optional H-MEM retrieval for semantic search.
        """
        self.session_memory = session_memory
        self.episode_store = episode_store
        self.memory_bridge = memory_bridge
        self.memory_router = memory_router
        self.hmem_retrieval = hmem_retrieval

    async def get_context_for_task(
        self,
        session_id: str,
        task: str,
        max_stores: int = 3,
    ) -> AssembledContext:
        """
        Assemble context relevant to a task.

        Args:
            session_id: Current session identifier.
            task: Task description or query.
            max_stores: Maximum number of stores to query.

        Returns:
            AssembledContext with data from relevant stores.
        """
        context = AssembledContext()

        # Determine which stores to query via router
        stores_to_query = await self._determine_stores(task, max_stores)
        context.routes_used = [s.value for s in stores_to_query]

        # Query each store
        for store_type in stores_to_query:
            await self._query_store(store_type, session_id, task, context)

        return context

    async def _determine_stores(
        self, task: str, max_stores: int
    ) -> List["MemoryStoreType"]:
        """
        Determine which stores to query based on task.

        Args:
            task: Task description.
            max_stores: Maximum stores to return.

        Returns:
            List of MemoryStoreType to query.
        """
        if self.memory_router is None:
            # Default stores if no router
            from gsd_rlm.memory.bridge.router import MemoryStoreType

            return [
                MemoryStoreType.SESSION,
                MemoryStoreType.HMEM_EPISODE,
                MemoryStoreType.BRIDGE_L2,
            ][:max_stores]

        # Use router to determine stores
        from gsd_rlm.memory.bridge.router import MemoryStoreType

        available_stores = self._get_available_stores()
        routes = await self.memory_router.route(task, available_stores)

        return self.memory_router.get_top_stores(routes, max_stores)

    def _get_available_stores(self) -> List["MemoryStoreType"]:
        """
        Get list of available stores based on configured components.

        Returns:
            List of MemoryStoreType that have configured backends.
        """
        from gsd_rlm.memory.bridge.router import MemoryStoreType

        stores: List[MemoryStoreType] = []

        if self.session_memory is not None:
            stores.append(MemoryStoreType.SESSION)

        if self.episode_store is not None:
            stores.append(MemoryStoreType.HMEM_EPISODE)

        if self.memory_bridge is not None:
            stores.extend(
                [
                    MemoryStoreType.BRIDGE_L0,
                    MemoryStoreType.BRIDGE_L1,
                    MemoryStoreType.BRIDGE_L2,
                    MemoryStoreType.BRIDGE_L3,
                ]
            )

        # Always include InfiniRetri as option (external system)
        stores.append(MemoryStoreType.INFINIRETRI)

        return stores

    async def _query_store(
        self,
        store_type: "MemoryStoreType",
        session_id: str,
        task: str,
        context: AssembledContext,
    ) -> None:
        """
        Query a specific store and add results to context.

        Args:
            store_type: Type of store to query.
            session_id: Current session identifier.
            task: Task description.
            context: AssembledContext to populate.
        """
        from gsd_rlm.memory.bridge.router import MemoryStoreType

        if store_type == MemoryStoreType.SESSION:
            await self._query_session(session_id, context)

        elif store_type == MemoryStoreType.HMEM_EPISODE:
            await self._query_episodes(session_id, task, context)

        elif store_type == MemoryStoreType.HMEM_TRACE:
            await self._query_traces(task, context)

        elif store_type == MemoryStoreType.BRIDGE_L0:
            await self._query_bridge_facts("L0", session_id, context)

        elif store_type == MemoryStoreType.BRIDGE_L1:
            await self._query_bridge_facts("L1", self._get_phase_id(), context)

        elif store_type == MemoryStoreType.BRIDGE_L2:
            await self._query_bridge_facts("L2", "project", context)

        elif store_type == MemoryStoreType.BRIDGE_L3:
            await self._query_bridge_facts("L3", "workspace", context)

    async def _query_session(self, session_id: str, context: AssembledContext) -> None:
        """
        Query session memory for current context.

        Args:
            session_id: Session identifier.
            context: AssembledContext to populate.
        """
        if self.session_memory is None:
            return

        try:
            # Get session data
            session_data = self._format_session_context(self.session_memory)
            context.session = session_data
        except Exception:
            context.session = {}

    async def _query_episodes(
        self, session_id: str, task: str, context: AssembledContext
    ) -> None:
        """
        Query H-MEM episodes for relevant experiences.

        Args:
            session_id: Session identifier.
            task: Task description.
            context: AssembledContext to populate.
        """
        if self.episode_store is None:
            return

        try:
            # Get recent episodes for this session
            episodes = self.episode_store.get_episodes_for_session(session_id)

            # Also try semantic search if available
            if self.hmem_retrieval is not None:
                search_results = await self.hmem_retrieval.search(
                    task, k=5, levels=["episode"]
                )
                # Get full content for search results
                for ep_id, score in search_results[:3]:
                    episode = self.episode_store.get(ep_id)
                    if episode and episode.episode_id not in [
                        e.get("episode_id") for e in context.recent_episodes
                    ]:
                        context.recent_episodes.append(
                            self._format_episode(episode, score)
                        )

            # Add session episodes (limit to avoid duplication)
            for episode in episodes[:5]:
                if episode.episode_id not in [
                    e.get("episode_id") for e in context.recent_episodes
                ]:
                    context.recent_episodes.append(self._format_episode(episode))
        except Exception:
            pass

    async def _query_traces(self, task: str, context: AssembledContext) -> None:
        """
        Query H-MEM traces for relevant work sessions.

        Args:
            task: Task description.
            context: AssembledContext to populate.
        """
        if self.hmem_retrieval is None:
            return

        try:
            search_results = await self.hmem_retrieval.search(
                task, k=3, levels=["trace"]
            )
            for trace_id, score in search_results:
                content, metadata = self.hmem_retrieval._get_trace_content(trace_id)
                if content:
                    context.relevant_traces.append(
                        {
                            "trace_id": trace_id,
                            "theme": metadata.get("theme", ""),
                            "summary": content,
                            "patterns": metadata.get("patterns", []),
                            "relevance_score": score,
                        }
                    )
        except Exception:
            pass

    async def _query_bridge_facts(
        self, level: str, scope_id: str, context: AssembledContext
    ) -> None:
        """
        Query Memory Bridge for facts at a specific level.

        Args:
            level: Bridge level (L0, L1, L2, L3).
            scope_id: Scope identifier.
            context: AssembledContext to populate.
        """
        if self.memory_bridge is None:
            return

        try:
            from gsd_rlm.memory.bridge.facts import BridgeLevel

            level_map = {
                "L0": BridgeLevel.L0_SESSION,
                "L1": BridgeLevel.L1_PHASE,
                "L2": BridgeLevel.L2_PROJECT,
                "L3": BridgeLevel.L3_WORKSPACE,
            }

            bridge_level = level_map.get(level)
            if bridge_level is None:
                return

            facts = self.memory_bridge.get_facts(bridge_level, scope_id)
            formatted = self._format_bridge_facts(facts)

            if level not in context.bridge_facts:
                context.bridge_facts[level] = {}

            context.bridge_facts[level].update(formatted)

            # Also update project_facts for L2
            if level == "L2":
                context.project_facts.update(formatted)
        except Exception:
            pass

    def _format_session_context(self, session: "FileSessionMemory") -> Dict[str, Any]:
        """
        Format session memory for context.

        Args:
            session: FileSessionMemory instance.

        Returns:
            Formatted session data.
        """
        result: Dict[str, Any] = {
            "messages": [],
            "recent_tasks": [],
        }

        try:
            # Get session history
            if hasattr(session, "get_history"):
                history = session.get_history()
                for entry in history[-10:]:
                    if isinstance(entry, dict):
                        role = entry.get("role", "unknown")
                        content = entry.get("content", "")
                        result["messages"].append(f"[{role}] {content[:100]}")
                    else:
                        result["messages"].append(str(entry)[:100])

            # Get recent tasks
            if hasattr(session, "get_recent_tasks"):
                tasks = session.get_recent_tasks()
                result["recent_tasks"] = [str(t) for t in tasks[-5:]]
        except Exception:
            pass

        return result

    def _format_episode(
        self, episode: Any, relevance_score: float = 1.0
    ) -> Dict[str, Any]:
        """
        Format an episode for context.

        Args:
            episode: Episode instance.
            relevance_score: Search relevance score.

        Returns:
            Formatted episode data.
        """
        return {
            "episode_id": episode.episode_id,
            "action": episode.action,
            "outcome": episode.outcome,
            "success": episode.success,
            "episode_type": episode.episode_type.value
            if hasattr(episode.episode_type, "value")
            else str(episode.episode_type),
            "tags": episode.tags,
            "relevance_score": relevance_score,
        }

    def _format_bridge_facts(self, facts: Dict[str, Any]) -> Dict[str, Any]:
        """
        Format bridge facts for context.

        Args:
            facts: Dictionary of BridgeFact instances.

        Returns:
            Formatted facts as key-value pairs.
        """
        result: Dict[str, Any] = {}
        for key, fact in facts.items():
            if hasattr(fact, "value"):
                result[key] = fact.value
            else:
                result[key] = fact
        return result

    def _get_phase_id(self) -> str:
        """
        Get current phase ID from session or environment.

        Returns:
            Phase identifier string.
        """
        # Try to get from session memory
        if self.session_memory is not None:
            try:
                if hasattr(self.session_memory, "get"):
                    phase = self.session_memory.get("phase_id")
                    if phase:
                        return phase
            except Exception:
                pass

        # Default phase ID
        return "03-memory-systems-infiniretri"
