"""
Sequence Reconstruction
=======================

Reconstructs event sequences from episodic memory.

Provides:
- Causal chain reconstruction
- Event ordering
- Gap detection
- Narrative generation
"""

from __future__ import annotations

import logging
from dataclasses import dataclass, field
from datetime import datetime, timedelta
from enum import Enum
from typing import Any, Callable, Dict, List, Optional, Set

logger = logging.getLogger(__name__)


class SequenceType(Enum):
    """Types of sequences that can be reconstructed."""
    TEMPORAL = "temporal"         # Time-ordered sequence
    CAUSAL = "causal"             # Cause-effect chain
    NARRATIVE = "narrative"       # Story-like sequence
    PROCEDURAL = "procedural"     # Step-by-step procedure


@dataclass
class SequenceNode:
    """A node in a reconstructed sequence."""
    episode_id: str
    position: int
    timestamp: Optional[datetime] = None
    content_preview: Optional[str] = None
    trust: float = 1.0
    role: Optional[str] = None  # "start", "middle", "end", "milestone"
    metadata: Dict[str, Any] = field(default_factory=dict)

    def to_dict(self) -> Dict[str, Any]:
        return {
            "episode_id": self.episode_id,
            "position": self.position,
            "timestamp": self.timestamp.isoformat() if self.timestamp else None,
            "content_preview": self.content_preview,
            "trust": self.trust,
            "role": self.role,
            "metadata": self.metadata,
        }


@dataclass
class SequenceGap:
    """Represents a detected gap in a sequence."""
    position: int
    before_episode_id: str
    after_episode_id: str
    time_gap: Optional[timedelta] = None
    estimated_missing: int = 0
    confidence: float = 0.5

    def to_dict(self) -> Dict[str, Any]:
        return {
            "position": self.position,
            "before_episode_id": self.before_episode_id,
            "after_episode_id": self.after_episode_id,
            "time_gap_seconds": self.time_gap.total_seconds() if self.time_gap else None,
            "estimated_missing": self.estimated_missing,
            "confidence": self.confidence,
        }


@dataclass
class ReconstructedSequence:
    """A fully reconstructed sequence of events."""
    sequence_type: SequenceType
    nodes: List[SequenceNode]
    gaps: List[SequenceGap] = field(default_factory=list)
    total_duration: Optional[timedelta] = None
    completeness: float = 1.0  # 0-1, how complete the sequence appears
    confidence: float = 1.0
    query: Optional[str] = None
    metadata: Dict[str, Any] = field(default_factory=dict)

    @property
    def length(self) -> int:
        return len(self.nodes)

    @property
    def start_node(self) -> Optional[SequenceNode]:
        return self.nodes[0] if self.nodes else None

    @property
    def end_node(self) -> Optional[SequenceNode]:
        return self.nodes[-1] if self.nodes else None

    def to_dict(self) -> Dict[str, Any]:
        return {
            "sequence_type": self.sequence_type.value,
            "nodes": [n.to_dict() for n in self.nodes],
            "gaps": [g.to_dict() for g in self.gaps],
            "length": self.length,
            "total_duration_seconds": self.total_duration.total_seconds() if self.total_duration else None,
            "completeness": self.completeness,
            "confidence": self.confidence,
            "query": self.query,
        }

    def to_narrative(self, connector: str = " → ") -> str:
        """Generate a simple narrative string."""
        if not self.nodes:
            return ""

        parts = []
        for node in self.nodes:
            if node.content_preview:
                parts.append(node.content_preview)
            else:
                parts.append(f"[{node.episode_id}]")

        return connector.join(parts)


class SequenceReconstructor:
    """
    Reconstructs event sequences from episodic memory.

    Supports multiple reconstruction strategies:
    - Temporal: Pure time-based ordering
    - Causal: Follow cause-effect edges
    - Semantic: Group by semantic similarity
    """

    # Edge types that indicate causal relationships
    CAUSAL_EDGE_TYPES = {
        "causes", "leads_to", "results_in", "triggers",
        "enables", "prevents", "follows", "precedes",
    }

    def __init__(
        self,
        gap_threshold: timedelta = timedelta(hours=4),
        min_sequence_length: int = 2,
        max_sequence_length: int = 100,
    ):
        """
        Initialize the reconstructor.

        Args:
            gap_threshold: Time gap threshold for detecting missing events
            min_sequence_length: Minimum nodes for a valid sequence
            max_sequence_length: Maximum nodes to include
        """
        self.gap_threshold = gap_threshold
        self.min_sequence_length = min_sequence_length
        self.max_sequence_length = max_sequence_length

    def reconstruct_temporal(
        self,
        episodes: List[Any],
        start_time: Optional[datetime] = None,
        end_time: Optional[datetime] = None,
        query: Optional[str] = None,
    ) -> ReconstructedSequence:
        """
        Reconstruct a purely temporal sequence.

        Args:
            episodes: List of episodes to order
            start_time: Optional start time filter
            end_time: Optional end time filter
            query: Optional query description

        Returns:
            ReconstructedSequence with time-ordered nodes
        """
        # Filter and sort by timestamp
        timed_episodes = []
        for ep in episodes:
            ts = getattr(ep, 'tocc', None)
            if ts is None:
                meta = getattr(ep, 'meta', {}) or {}
                ts = meta.get('timestamp')
                if isinstance(ts, str):
                    ts = datetime.fromisoformat(ts)

            if ts is None:
                continue

            if start_time and ts < start_time:
                continue
            if end_time and ts > end_time:
                continue

            timed_episodes.append((ts, ep))

        # Sort by time
        timed_episodes.sort(key=lambda x: x[0])

        # Limit
        timed_episodes = timed_episodes[:self.max_sequence_length]

        # Build nodes
        nodes = []
        for i, (ts, ep) in enumerate(timed_episodes):
            role = None
            if i == 0:
                role = "start"
            elif i == len(timed_episodes) - 1:
                role = "end"

            nodes.append(SequenceNode(
                episode_id=ep.id,
                position=i,
                timestamp=ts,
                content_preview=self._get_preview(ep),
                trust=getattr(ep, 'trust', 1.0),
                role=role,
            ))

        # Detect gaps
        gaps = self._detect_temporal_gaps(nodes)

        # Calculate duration
        duration = None
        if len(nodes) >= 2:
            duration = nodes[-1].timestamp - nodes[0].timestamp

        # Estimate completeness
        completeness = self._estimate_completeness(nodes, gaps)

        return ReconstructedSequence(
            sequence_type=SequenceType.TEMPORAL,
            nodes=nodes,
            gaps=gaps,
            total_duration=duration,
            completeness=completeness,
            confidence=1.0,  # High confidence for pure temporal
            query=query,
        )

    def reconstruct_causal(
        self,
        start_episode: Any,
        episodes: List[Any],
        graph_lookup: Optional[Callable[[str], Any]] = None,
        max_hops: int = 10,
    ) -> ReconstructedSequence:
        """
        Reconstruct a causal sequence starting from an episode.

        Follows causal edges (causes, leads_to, etc.) through the graph.

        Args:
            start_episode: Starting episode
            episodes: Pool of episodes
            graph_lookup: Function to get episode graph
            max_hops: Maximum causal hops to follow

        Returns:
            ReconstructedSequence with causal ordering
        """
        episode_map = {ep.id: ep for ep in episodes}
        visited: Set[str] = set()
        sequence: List[Any] = []

        # BFS through causal edges
        queue = [start_episode]

        while queue and len(sequence) < self.max_sequence_length:
            current = queue.pop(0)
            if current.id in visited:
                continue

            visited.add(current.id)
            sequence.append(current)

            if len(sequence) >= max_hops:
                break

            # Find causal successors
            graph = getattr(current, 'G', None)
            if graph is None and graph_lookup:
                graph = graph_lookup(current.id)

            if graph is not None:
                for _, target, data in graph.edges(data=True):
                    rel = (data.get('relation') or data.get('label') or '').lower()
                    if rel in self.CAUSAL_EDGE_TYPES:
                        # Target might be an episode ID
                        if target in episode_map and target not in visited:
                            queue.append(episode_map[target])

        # Build nodes
        nodes = []
        for i, ep in enumerate(sequence):
            role = "start" if i == 0 else ("end" if i == len(sequence) - 1 else None)
            nodes.append(SequenceNode(
                episode_id=ep.id,
                position=i,
                timestamp=getattr(ep, 'tocc', None),
                content_preview=self._get_preview(ep),
                trust=getattr(ep, 'trust', 1.0),
                role=role,
            ))

        return ReconstructedSequence(
            sequence_type=SequenceType.CAUSAL,
            nodes=nodes,
            gaps=[],  # Causal chains don't have "gaps" in the same sense
            total_duration=None,
            completeness=1.0 if nodes else 0.0,
            confidence=0.8,  # Lower confidence for inferred causality
            query=f"Causal chain from {start_episode.id}",
        )

    def reconstruct_narrative(
        self,
        episodes: List[Any],
        topic_vector: Optional[Any] = None,
        similarity_fn: Optional[Callable] = None,
        query: Optional[str] = None,
    ) -> ReconstructedSequence:
        """
        Reconstruct a narrative sequence.

        Combines temporal ordering with semantic coherence.

        Args:
            episodes: Episodes to form narrative from
            topic_vector: Optional topic embedding for relevance filtering
            similarity_fn: Optional similarity function
            query: Query description

        Returns:
            ReconstructedSequence as a narrative
        """
        import numpy as np

        # Score episodes by relevance if topic provided
        if topic_vector is not None:
            scored = []
            for ep in episodes:
                embedding = getattr(ep, 'e', None)
                if embedding is not None:
                    if similarity_fn:
                        sim = similarity_fn(topic_vector, embedding)
                    else:
                        norm_t = np.linalg.norm(topic_vector)
                        norm_e = np.linalg.norm(embedding)
                        sim = float(np.dot(topic_vector, embedding) / (norm_t * norm_e)) if norm_t and norm_e else 0
                    scored.append((sim, ep))
                else:
                    scored.append((0.0, ep))

            # Filter to relevant episodes
            scored.sort(key=lambda x: x[0], reverse=True)
            threshold = 0.3
            relevant = [ep for sim, ep in scored if sim >= threshold]
        else:
            relevant = episodes

        # Sort remaining by time to form narrative
        timed = []
        for ep in relevant:
            ts = getattr(ep, 'tocc', None)
            if ts:
                timed.append((ts, ep))

        timed.sort(key=lambda x: x[0])
        timed = timed[:self.max_sequence_length]

        # Build nodes with narrative roles
        nodes = []
        for i, (ts, ep) in enumerate(timed):
            # Assign narrative roles
            pct = i / max(len(timed) - 1, 1)
            if pct < 0.1:
                role = "start"
            elif pct > 0.9:
                role = "end"
            elif 0.4 <= pct <= 0.6:
                role = "milestone"
            else:
                role = "middle"

            nodes.append(SequenceNode(
                episode_id=ep.id,
                position=i,
                timestamp=ts,
                content_preview=self._get_preview(ep),
                trust=getattr(ep, 'trust', 1.0),
                role=role,
            ))

        # Detect gaps
        gaps = self._detect_temporal_gaps(nodes)

        # Duration
        duration = None
        if len(nodes) >= 2:
            duration = nodes[-1].timestamp - nodes[0].timestamp

        return ReconstructedSequence(
            sequence_type=SequenceType.NARRATIVE,
            nodes=nodes,
            gaps=gaps,
            total_duration=duration,
            completeness=self._estimate_completeness(nodes, gaps),
            confidence=0.7,  # Lower confidence for narrative inference
            query=query,
        )

    def find_milestones(
        self,
        sequence: ReconstructedSequence,
        milestone_count: int = 5,
    ) -> List[SequenceNode]:
        """
        Find key milestones in a sequence.

        Uses a combination of:
        - Temporal spacing (evenly distributed)
        - Trust scores (high trust = important)
        - Explicit roles

        Args:
            sequence: The sequence to analyze
            milestone_count: Number of milestones to identify

        Returns:
            List of milestone nodes
        """
        if not sequence.nodes:
            return []

        if len(sequence.nodes) <= milestone_count:
            return sequence.nodes

        # Always include start and end
        milestones = [sequence.nodes[0], sequence.nodes[-1]]

        # Score remaining nodes
        remaining = sequence.nodes[1:-1]
        scored = []

        for i, node in enumerate(remaining):
            # Position score (prefer evenly spaced)
            position_pct = (i + 1) / (len(sequence.nodes) - 1)
            ideal_positions = [j / (milestone_count - 1) for j in range(1, milestone_count - 1)]
            position_score = 1 - min(abs(position_pct - ideal) for ideal in ideal_positions)

            # Trust score
            trust_score = node.trust

            # Combined score
            score = 0.6 * position_score + 0.4 * trust_score
            scored.append((score, node))

        # Select top milestones
        scored.sort(key=lambda x: x[0], reverse=True)
        for _, node in scored[:milestone_count - 2]:
            milestones.append(node)

        # Sort by position
        milestones.sort(key=lambda n: n.position)

        return milestones

    def _get_preview(self, episode: Any) -> Optional[str]:
        """Extract content preview from episode."""
        meta = getattr(episode, 'meta', {}) or {}
        preview = (
            meta.get('text') or
            meta.get('content') or
            meta.get('content_preview')
        )
        if preview and len(preview) > 100:
            preview = preview[:100] + "..."
        return preview

    def _detect_temporal_gaps(
        self,
        nodes: List[SequenceNode],
    ) -> List[SequenceGap]:
        """Detect temporal gaps in a sequence."""
        gaps = []

        for i in range(len(nodes) - 1):
            current = nodes[i]
            next_node = nodes[i + 1]

            if current.timestamp is None or next_node.timestamp is None:
                continue

            time_gap = next_node.timestamp - current.timestamp

            if time_gap > self.gap_threshold:
                # Estimate missing events based on average interval
                avg_interval = self.gap_threshold / 2
                estimated_missing = int(time_gap / avg_interval) - 1

                gaps.append(SequenceGap(
                    position=i + 1,
                    before_episode_id=current.episode_id,
                    after_episode_id=next_node.episode_id,
                    time_gap=time_gap,
                    estimated_missing=max(0, estimated_missing),
                    confidence=0.6,
                ))

        return gaps

    def _estimate_completeness(
        self,
        nodes: List[SequenceNode],
        gaps: List[SequenceGap],
    ) -> float:
        """Estimate sequence completeness (0-1)."""
        if not nodes:
            return 0.0

        if not gaps:
            return 1.0

        total_estimated_missing = sum(g.estimated_missing for g in gaps)
        completeness = len(nodes) / (len(nodes) + total_estimated_missing)

        return min(1.0, max(0.0, completeness))

# ------------------------------------------------------------------ #
# API Helpers (Exported for API Routes)
# ------------------------------------------------------------------ #


def reconstruct_sequence(
    memory: Any,  # generic to avoid circular import with Memory
    start_episode_id: Optional[str] = None,
    end_episode_id: Optional[str] = None,
    limit: int = 100,
    namespace: Optional[str] = None,
) -> List[Any]:
    """
    Reconstruct a sequence of episodes from memory.

    Args:
        memory: Memory instance
        start_episode_id: Optional start episode ID
        end_episode_id: Optional end episode ID
        limit: Max episodes
        namespace: Namespace to use

    Returns:
        List of Episode objects
    """
    reconstructor = SequenceReconstructor(max_sequence_length=limit)

    with memory.use_namespace(namespace):
        # Determine strategy based on inputs
        if start_episode_id and end_episode_id:
            # We want a path/sequence between two points
            # This is hard to do without temporal index or graph search
            # For now, simplistic temporal range filtering
            start_ep = memory._get_episode_by_id(start_episode_id)
            end_ep = memory._get_episode_by_id(end_episode_id)

            if not start_ep or not end_ep:
                return []

            # Use temporal reconstruction on the range
            seq = reconstructor.reconstruct_temporal(
                memory.episodes,
                start_time=getattr(start_ep, 'tocc', None),
                end_time=getattr(end_ep, 'tocc', None),
            )

            # Map nodes back to episodes
            episodes = []
            for node in seq.nodes:
                ep = memory._get_episode_by_id(node.episode_id)
                if ep:
                    episodes.append(ep)
            return episodes

        elif start_episode_id:
            # Causal chain from start
            start_ep = memory._get_episode_by_id(start_episode_id)
            if not start_ep:
                return []

            seq = reconstructor.reconstruct_causal(
                start_ep,
                memory.episodes,
                graph_lookup=lambda eid: getattr(memory._get_episode_by_id(eid), 'G', None)
            )
            # Map nodes back to episodes
            episodes = []
            for node in seq.nodes:
                ep = memory._get_episode_by_id(node.episode_id)
                if ep:
                    episodes.append(ep)
            return episodes

        else:
            # Fallback to simple recent temporal
            seq = reconstructor.reconstruct_temporal(memory.episodes)
            # Map nodes back to episodes
            episodes = []
            for node in seq.nodes:
                ep = memory._get_episode_by_id(node.episode_id)
                if ep:
                    episodes.append(ep)
            return episodes
