"""
Session Resumption for GSD-RLM Memory System.

Provides full context restoration for session resumption, combining
H-MEM episodes, Memory Bridge facts, and recent traces into a unified
continuation context for agents.

Key components:
- ResumptionContext: Dataclass containing all restored context
- SessionResumption: Restores full context from all memory stores
"""

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

if TYPE_CHECKING:
    from gsd_rlm.session.memory import FileSessionMemory, SessionState
    from gsd_rlm.memory.hmem.store import EpisodeStore
    from gsd_rlm.memory.bridge.store import MemoryBridge
    from gsd_rlm.memory.bridge.router import MemoryRouter
    from gsd_rlm.memory.hmem.trace import Trace


def _utcnow_iso() -> str:
    """Get current UTC time as ISO string (timezone-aware)."""
    return datetime.now(timezone.utc).isoformat()


@dataclass
class ResumptionContext:
    """
    Full context for session resumption.

    Contains all the information needed to restore an agent's context
    when resuming a previous session.

    Attributes:
        session_id: Unique identifier of the session being resumed.
        session_data: Raw session data from FileSessionMemory.
        recent_episodes: Recent H-MEM episodes from the session.
        relevant_traces: Consolidated traces relevant to the session.
        project_facts: Project-level facts from Memory Bridge (L2).
        phase_facts: Phase-level facts from Memory Bridge (L1).
        last_task: The last task that was being worked on.
        next_suggested_actions: AI-suggested next actions based on context.
        restored_at: Timestamp when context was restored.
    """

    session_id: str
    session_data: Dict[str, Any] = field(default_factory=dict)
    recent_episodes: List[Dict[str, Any]] = field(default_factory=list)
    relevant_traces: List[Dict[str, Any]] = field(default_factory=list)
    project_facts: Dict[str, Any] = field(default_factory=dict)
    phase_facts: Dict[str, Any] = field(default_factory=dict)
    last_task: Optional[str] = None
    next_suggested_actions: List[str] = field(default_factory=list)
    restored_at: str = field(default_factory=_utcnow_iso)

    def to_dict(self) -> Dict[str, Any]:
        """
        Convert to dictionary for serialization.

        Returns:
            Dictionary representation of the resumption context.
        """
        return {
            "session_id": self.session_id,
            "session_data": self.session_data,
            "recent_episodes": self.recent_episodes,
            "relevant_traces": self.relevant_traces,
            "project_facts": self.project_facts,
            "phase_facts": self.phase_facts,
            "last_task": self.last_task,
            "next_suggested_actions": self.next_suggested_actions,
            "restored_at": self.restored_at,
        }

    @classmethod
    def from_dict(cls, data: Dict[str, Any]) -> "ResumptionContext":
        """
        Create ResumptionContext from dictionary.

        Args:
            data: Dictionary with resumption context data.

        Returns:
            ResumptionContext instance.
        """
        return cls(
            session_id=data.get("session_id", ""),
            session_data=data.get("session_data", {}),
            recent_episodes=data.get("recent_episodes", []),
            relevant_traces=data.get("relevant_traces", []),
            project_facts=data.get("project_facts", {}),
            phase_facts=data.get("phase_facts", {}),
            last_task=data.get("last_task"),
            next_suggested_actions=data.get("next_suggested_actions", []),
            restored_at=data.get("restored_at", _utcnow_iso()),
        )


class SessionResumption:
    """
    Restores full context for session resumption (MEM-12).

    Combines data from multiple memory stores to provide a complete
    restoration of agent context when resuming a previous session.

    Flow:
    1. Load session_data from FileSessionMemory
    2. Get recent episodes from EpisodeStore
    3. Get project facts from MemoryBridge (L2_PROJECT)
    4. Get phase facts from MemoryBridge (L1_PHASE)
    5. Determine last_task from session_data
    6. Generate next_suggested_actions using MemoryRouter
    7. Return ResumptionContext

    Attributes:
        session_memory: FileSessionMemory for session data.
        episode_store: EpisodeStore for H-MEM episodes.
        memory_bridge: MemoryBridge for persistent facts.
        memory_router: MemoryRouter for context routing.

    Example:
        >>> from gsd_rlm.session.memory import FileSessionMemory
        >>> from gsd_rlm.memory.hmem.store import EpisodeStore
        >>> from gsd_rlm.memory.bridge.store import MemoryBridge
        >>> resumption = SessionResumption(
        ...     session_memory=FileSessionMemory(Path("sessions")),
        ...     episode_store=EpisodeStore("memory/hmem.db"),
        ...     memory_bridge=MemoryBridge(Path(".")),
        ... )
        >>> context = await resumption.resume("session-001")
    """

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

        Args:
            session_memory: FileSessionMemory for session data.
            episode_store: EpisodeStore for H-MEM episodes.
            memory_bridge: MemoryBridge for persistent facts.
            memory_router: MemoryRouter for context routing.
        """
        self.session_memory = session_memory
        self.episode_store = episode_store
        self.memory_bridge = memory_bridge
        self.memory_router = memory_router

    async def resume(self, session_id: str) -> ResumptionContext:
        """
        Restore full context for a session.

        Args:
            session_id: Session identifier to resume.

        Returns:
            ResumptionContext with all restored data.

        Raises:
            ValueError: If session not found.
        """
        # 1. Load session data
        session_data = await self._load_session_data(session_id)

        # 2. Get recent episodes
        recent_episodes = await self._get_recent_episodes(session_id)

        # 3. Get relevant traces
        relevant_traces = await self._get_relevant_traces(session_data)

        # 4. Get project facts (L2_PROJECT)
        project_facts = await self._get_project_facts()

        # 5. Get phase facts (L1_PHASE)
        phase_id = self._get_current_phase(session_data)
        phase_facts = await self._get_phase_facts(phase_id)

        # 6. Determine last task
        last_task = self._get_last_task(session_data)

        # 7. Generate next suggested actions
        next_actions = await self._suggest_next_actions(
            session_data, recent_episodes, project_facts
        )

        return ResumptionContext(
            session_id=session_id,
            session_data=session_data,
            recent_episodes=recent_episodes,
            relevant_traces=relevant_traces,
            project_facts=project_facts,
            phase_facts=phase_facts,
            last_task=last_task,
            next_suggested_actions=next_actions,
        )

    async def _load_session_data(self, session_id: str) -> Dict[str, Any]:
        """
        Load session data from FileSessionMemory.

        Args:
            session_id: Session identifier.

        Returns:
            Session data dictionary.

        Raises:
            ValueError: If session not found.
        """
        if self.session_memory is None:
            return {}

        session = self.session_memory.load(session_id)
        if session is None:
            raise ValueError(f"Session not found: {session_id}")

        return session.to_dict()

    async def _get_recent_episodes(self, session_id: str) -> List[Dict[str, Any]]:
        """
        Get recent episodes for the session.

        Args:
            session_id: Session identifier.

        Returns:
            List of episode dictionaries.
        """
        if self.episode_store is None:
            return []

        try:
            episodes = self.episode_store.get_episodes_for_session(session_id)
            return [self._format_episode(ep) for ep in episodes[-20:]]  # Last 20
        except Exception:
            return []

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

        Args:
            episode: Episode instance.

        Returns:
            Formatted episode dictionary.
        """
        return {
            "episode_id": episode.episode_id,
            "episode_type": episode.episode_type.value
            if hasattr(episode.episode_type, "value")
            else str(episode.episode_type),
            "context": episode.context[:500] if episode.context else "",
            "action": episode.action,
            "outcome": episode.outcome[:500] if episode.outcome else "",
            "success": episode.success,
            "timestamp": episode.timestamp,
            "tags": episode.tags,
        }

    async def _get_relevant_traces(
        self, session_data: Dict[str, Any]
    ) -> List[Dict[str, Any]]:
        """
        Get relevant traces based on session context.

        Args:
            session_data: Session data for context.

        Returns:
            List of trace dictionaries.
        """
        # Extract keywords from session for trace matching
        keywords = self._extract_keywords_from_session(session_data)

        # For now, return empty list as trace retrieval requires
        # H-MEM retrieval integration
        # TODO: Integrate with HMEMRetrieval for semantic trace search
        return []

    async def _get_project_facts(self) -> Dict[str, Any]:
        """
        Get project-level facts from Memory Bridge (L2_PROJECT).

        Returns:
            Dictionary of project facts.
        """
        if self.memory_bridge is None:
            return {}

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

            facts = self.memory_bridge.get_facts(BridgeLevel.L2_PROJECT, "project")
            return self._format_bridge_facts(facts)
        except Exception:
            return {}

    async def _get_phase_facts(self, phase_id: str) -> Dict[str, Any]:
        """
        Get phase-level facts from Memory Bridge (L1_PHASE).

        Args:
            phase_id: Phase identifier.

        Returns:
            Dictionary of phase facts.
        """
        if self.memory_bridge is None:
            return {}

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

            facts = self.memory_bridge.get_facts(BridgeLevel.L1_PHASE, phase_id)
            return self._format_bridge_facts(facts)
        except Exception:
            return {}

    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_current_phase(self, session_data: Dict[str, Any]) -> str:
        """
        Determine current phase from session data.

        Args:
            session_data: Session data dictionary.

        Returns:
            Phase identifier string.
        """
        # Try to get from metadata
        metadata = session_data.get("metadata", {})
        if "phase_id" in metadata:
            return metadata["phase_id"]

        # Try to infer from current task
        current_task = session_data.get("current_task", "")
        if current_task:
            # Look for phase pattern like "03-" or "phase 3"
            import re

            match = re.search(r"(\d+)-", current_task)
            if match:
                return f"{match.group(1)}-phase"

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

    def _get_last_task(self, session_data: Dict[str, Any]) -> Optional[str]:
        """
        Determine the last task from session data.

        Args:
            session_data: Session data dictionary.

        Returns:
            Last task description or None.
        """
        # Check current_task first
        current_task = session_data.get("current_task")
        if current_task:
            return current_task

        # Check task outputs for most recent
        task_outputs = session_data.get("task_outputs", [])
        if task_outputs:
            last_output = task_outputs[-1]
            return last_output.get("task")

        # Check messages for recent task mentions
        messages = session_data.get("messages", [])
        for msg in reversed(messages[-10:]):
            content = msg.get("content", "")
            if "task" in content.lower():
                return content[:200]  # Truncate long messages

        return None

    async def _suggest_next_actions(
        self,
        session_data: Dict[str, Any],
        recent_episodes: List[Dict[str, Any]],
        project_facts: Dict[str, Any],
    ) -> List[str]:
        """
        Generate suggested next actions based on context.

        Args:
            session_data: Session data dictionary.
            recent_episodes: Recent episode data.
            project_facts: Project-level facts.

        Returns:
            List of suggested action strings.
        """
        suggestions: List[str] = []

        # Analyze recent episodes for incomplete work
        for episode in recent_episodes[-5:]:
            if not episode.get("success"):
                suggestions.append(
                    f"Retry failed action: {episode.get('action', 'unknown')}"
                )

        # Check for pending tasks in session
        metadata = session_data.get("metadata", {})
        pending_tasks = metadata.get("pending_tasks", [])
        for task in pending_tasks[:3]:
            suggestions.append(f"Continue with: {task}")

        # Add context-aware suggestions based on project facts
        if "roadmap_phase" in project_facts:
            suggestions.append(
                f"Continue work on phase: {project_facts['roadmap_phase']}"
            )

        # Limit to 5 suggestions
        return suggestions[:5]

    def _extract_keywords_from_session(self, session_data: Dict[str, Any]) -> List[str]:
        """
        Extract keywords from session data for trace matching.

        Args:
            session_data: Session data dictionary.

        Returns:
            List of keywords.
        """
        keywords: set = set()

        # Extract from current task
        current_task = session_data.get("current_task", "")
        if current_task:
            # Simple keyword extraction (split on common delimiters)
            import re

            words = re.findall(r"\b[a-zA-Z]{3,}\b", current_task.lower())
            keywords.update(words[:10])

        # Extract from recent messages
        messages = session_data.get("messages", [])
        for msg in messages[-5:]:
            content = msg.get("content", "")
            words = re.findall(r"\b[a-zA-Z]{3,}\b", content.lower())
            keywords.update(words[:5])

        return list(keywords)[:20]

    def generate_resumption_prompt(self, context: ResumptionContext) -> str:
        """
        Generate a continuation prompt for the agent.

        Args:
            context: ResumptionContext with restored data.

        Returns:
            Formatted continuation prompt string.
        """
        lines: List[str] = [
            "# Session Resumption Context",
            "",
            f"**Session ID:** {context.session_id}",
            f"**Restored at:** {context.restored_at}",
            "",
        ]

        # Last task
        if context.last_task:
            lines.extend(
                [
                    "## Last Task",
                    context.last_task,
                    "",
                ]
            )

        # Recent episodes
        if context.recent_episodes:
            lines.append("## Recent Episodes")
            for ep in context.recent_episodes[-5:]:
                success_marker = "✓" if ep.get("success") else "✗"
                action = ep.get("action", "Unknown action")
                lines.append(f"- [{success_marker}] {action}")
            lines.append("")

        # Project facts
        if context.project_facts:
            lines.append("## Project Facts")
            for key, value in list(context.project_facts.items())[:10]:
                if isinstance(value, dict):
                    lines.append(f"- **{key}:** {value}")
                else:
                    lines.append(f"- **{key}:** {value}")
            lines.append("")

        # Phase facts
        if context.phase_facts:
            lines.append("## Phase Facts")
            for key, value in list(context.phase_facts.items())[:5]:
                if isinstance(value, dict):
                    lines.append(f"- **{key}:** {value}")
                else:
                    lines.append(f"- **{key}:** {value}")
            lines.append("")

        # Suggested next actions
        if context.next_suggested_actions:
            lines.append("## Suggested Next Actions")
            for action in context.next_suggested_actions:
                lines.append(f"- {action}")
            lines.append("")

        # Continuation prompt
        lines.extend(
            [
                "---",
                "",
                "## Continuation Prompt",
                "",
                "You are resuming a previous session. Based on the context above:",
            ]
        )

        if context.last_task:
            lines.append(f"- Your last task was: **{context.last_task}**")

        if context.next_suggested_actions:
            lines.append(
                f"- Consider continuing with: **{context.next_suggested_actions[0]}**"
            )

        lines.extend(
            [
                "",
                "Please continue from where you left off, taking into account the "
                "recent episodes and project context.",
            ]
        )

        return "\n".join(lines)
