"""
Temporal Reasoning
==================

Provides temporal queries and reasoning over episodic memory.

Supports:
- Before/after queries
- Time range retrieval
- Sequence reconstruction
- Temporal decay adjustment
"""

from __future__ import annotations

import bisect
from dataclasses import dataclass, field
from datetime import datetime, timedelta
from typing import Any, Callable, Dict, List, Optional, Tuple, Union

import numpy as np


@dataclass
class TemporalQuery:
    """A temporal query specification."""
    query_text: Optional[str] = None
    query_vector: Optional[np.ndarray] = None
    start_time: Optional[datetime] = None
    end_time: Optional[datetime] = None
    reference_episode_id: Optional[str] = None
    relation: str = "around"  # "before", "after", "around", "between"
    limit: int = 10
    namespace: Optional[str] = None


@dataclass
class TemporalEvent:
    """An episode with temporal context."""
    episode_id: str
    timestamp: datetime
    content_preview: Optional[str] = None
    trust: float = 1.0
    metadata: Dict[str, Any] = field(default_factory=dict)
    temporal_distance: Optional[timedelta] = None  # Distance from query reference

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


@dataclass
class TemporalSequence:
    """A reconstructed sequence of events."""
    events: List[TemporalEvent]
    start_event: Optional[TemporalEvent] = None
    end_event: Optional[TemporalEvent] = None
    total_duration: Optional[timedelta] = None
    gap_count: int = 0  # Number of potential missing events

    def to_dict(self) -> Dict[str, Any]:
        return {
            "events": [e.to_dict() for e in self.events],
            "event_count": len(self.events),
            "start_time": self.start_event.timestamp.isoformat() if self.start_event else None,
            "end_time": self.end_event.timestamp.isoformat() if self.end_event else None,
            "total_duration_seconds": (
                self.total_duration.total_seconds() if self.total_duration else None
            ),
            "gap_count": self.gap_count,
        }


class TemporalIndex:
    """
    Index for efficient temporal queries.

    Maintains a sorted list of (timestamp, episode_id) pairs
    for fast before/after/range queries.
    """

    def __init__(self) -> None:
        """Initialize the temporal index."""
        # Sorted list of (timestamp, episode_id) tuples
        self._timeline: List[Tuple[datetime, str]] = []
        # Episode ID to timestamp lookup
        self._episode_times: Dict[str, datetime] = {}

    def add(self, episode_id: str, timestamp: datetime) -> None:
        """Add an episode to the temporal index."""
        if episode_id in self._episode_times:
            # Update existing
            self.remove(episode_id)

        entry = (timestamp, episode_id)
        bisect.insort(self._timeline, entry)
        self._episode_times[episode_id] = timestamp

    def remove(self, episode_id: str) -> bool:
        """Remove an episode from the index."""
        if episode_id not in self._episode_times:
            return False

        timestamp = self._episode_times[episode_id]
        entry = (timestamp, episode_id)

        try:
            self._timeline.remove(entry)
        except ValueError:
            pass

        del self._episode_times[episode_id]
        return True

    def get_timestamp(self, episode_id: str) -> Optional[datetime]:
        """Get the timestamp of an episode."""
        return self._episode_times.get(episode_id)

    def query_before(
        self,
        reference: Union[datetime, str],
        limit: int = 10,
    ) -> List[Tuple[datetime, str]]:
        """
        Get episodes before a reference point.

        Args:
            reference: Either a datetime or an episode ID
            limit: Maximum results to return

        Returns:
            List of (timestamp, episode_id) tuples, newest first
        """
        if isinstance(reference, str):
            ref_time = self._episode_times.get(reference)
            if ref_time is None:
                return []
        else:
            ref_time = reference

        # Find insertion point (where ref_time would be inserted)
        idx = bisect.bisect_left(self._timeline, (ref_time, ""))

        # Get items before this index
        results = self._timeline[max(0, idx - limit):idx]
        return list(reversed(results))  # Newest first

    def query_after(
        self,
        reference: Union[datetime, str],
        limit: int = 10,
    ) -> List[Tuple[datetime, str]]:
        """
        Get episodes after a reference point.

        Args:
            reference: Either a datetime or an episode ID
            limit: Maximum results to return

        Returns:
            List of (timestamp, episode_id) tuples, oldest first
        """
        if isinstance(reference, str):
            ref_time = self._episode_times.get(reference)
            if ref_time is None:
                return []
        else:
            ref_time = reference

        # Find insertion point
        idx = bisect.bisect_right(self._timeline, (ref_time, "\xff"))

        # Get items after this index
        return self._timeline[idx:idx + limit]

    def query_range(
        self,
        start: datetime,
        end: datetime,
        limit: Optional[int] = None,
    ) -> List[Tuple[datetime, str]]:
        """
        Get episodes within a time range.

        Args:
            start: Start of time range (inclusive)
            end: End of time range (inclusive)
            limit: Optional maximum results

        Returns:
            List of (timestamp, episode_id) tuples, oldest first
        """
        start_idx = bisect.bisect_left(self._timeline, (start, ""))
        end_idx = bisect.bisect_right(self._timeline, (end, "\xff"))

        results = self._timeline[start_idx:end_idx]

        if limit:
            results = results[:limit]

        return results

    def query_around(
        self,
        reference: Union[datetime, str],
        window: timedelta,
        limit: int = 10,
    ) -> List[Tuple[datetime, str]]:
        """
        Get episodes around a reference point within a time window.

        Args:
            reference: Either a datetime or an episode ID
            window: Time window (both before and after)
            limit: Maximum results to return

        Returns:
            List of (timestamp, episode_id) tuples, sorted by distance from reference
        """
        if isinstance(reference, str):
            ref_time = self._episode_times.get(reference)
            if ref_time is None:
                return []
        else:
            ref_time = reference

        start = ref_time - window
        end = ref_time + window

        results = self.query_range(start, end)

        # Sort by distance from reference
        results_with_distance = [
            (abs((ts - ref_time).total_seconds()), ts, ep_id)
            for ts, ep_id in results
        ]
        results_with_distance.sort(key=lambda x: x[0])

        return [(ts, ep_id) for _, ts, ep_id in results_with_distance[:limit]]

    def __len__(self) -> int:
        return len(self._timeline)


class TemporalReasoner:
    """
    Provides temporal reasoning capabilities over episodic memory.

    Integrates with Memory to answer temporal queries like:
    - "What happened before X?"
    - "What happened in January 2026?"
    - "Reconstruct the sequence from start to end"
    """

    def __init__(
        self,
        temporal_index: Optional[TemporalIndex] = None,
        recency_boost_factor: float = 1.2,
        max_gap_threshold: timedelta = timedelta(hours=24),
    ):
        """
        Initialize the temporal reasoner.

        Args:
            temporal_index: Optional pre-built temporal index
            recency_boost_factor: Boost factor for recent memories in hybrid scoring
            max_gap_threshold: Maximum gap before flagging missing events
        """
        self.index = temporal_index or TemporalIndex()
        self.recency_boost_factor = recency_boost_factor
        self.max_gap_threshold = max_gap_threshold

    def index_episode(self, episode: Any) -> None:
        """
        Add an episode to the temporal index.

        Args:
            episode: Episode to index
        """
        timestamp = getattr(episode, 'tocc', None)
        if timestamp is None:
            # Try to get from metadata
            meta = getattr(episode, 'meta', {}) or {}
            timestamp = meta.get('timestamp') or meta.get('created_at')
            if isinstance(timestamp, str):
                timestamp = datetime.fromisoformat(timestamp)

        if timestamp:
            self.index.add(episode.id, timestamp)

    def remove_episode(self, episode_id: str) -> bool:
        """Remove an episode from the temporal index."""
        return self.index.remove(episode_id)

    def recall_before(
        self,
        reference: Union[datetime, str, Any],  # datetime, episode_id, or Episode
        limit: int = 10,
        episode_lookup: Optional[Callable[[str], Any]] = None,
    ) -> List[TemporalEvent]:
        """
        Recall events that happened before a reference point.

        Args:
            reference: Reference point (datetime, episode_id, or Episode)
            limit: Maximum results
            episode_lookup: Function to look up episodes by ID

        Returns:
            List of TemporalEvent objects, newest first
        """
        ref = self._normalize_reference(reference)
        results = self.index.query_before(ref, limit)

        return self._to_temporal_events(results, ref, episode_lookup)

    def recall_after(
        self,
        reference: Union[datetime, str, Any],
        limit: int = 10,
        episode_lookup: Optional[Callable[[str], Any]] = None,
    ) -> List[TemporalEvent]:
        """
        Recall events that happened after a reference point.

        Args:
            reference: Reference point (datetime, episode_id, or Episode)
            limit: Maximum results
            episode_lookup: Function to look up episodes by ID

        Returns:
            List of TemporalEvent objects, oldest first
        """
        ref = self._normalize_reference(reference)
        results = self.index.query_after(ref, limit)

        return self._to_temporal_events(results, ref, episode_lookup)

    def recall_temporal(
        self,
        start: Union[datetime, str],
        end: Union[datetime, str],
        query_vector: Optional[np.ndarray] = None,
        limit: int = 20,
        episode_lookup: Optional[Callable[[str], Any]] = None,
        similarity_fn: Optional[Callable[[np.ndarray, np.ndarray], float]] = None,
    ) -> List[TemporalEvent]:
        """
        Recall events within a time range, optionally filtered by semantic similarity.

        Args:
            start: Start of time range
            end: End of time range
            query_vector: Optional query embedding for semantic filtering
            limit: Maximum results
            episode_lookup: Function to look up episodes by ID
            similarity_fn: Optional similarity function

        Returns:
            List of TemporalEvent objects
        """
        # Parse datetime strings
        if isinstance(start, str):
            start = datetime.fromisoformat(start)
        if isinstance(end, str):
            end = datetime.fromisoformat(end)

        results = self.index.query_range(start, end, limit=None)

        events = self._to_temporal_events(results, start, episode_lookup)

        # If query vector provided, re-rank by semantic similarity
        if query_vector is not None and episode_lookup is not None:
            events = self._rerank_by_similarity(
                events, query_vector, episode_lookup, similarity_fn
            )

        return events[:limit]

    def reconstruct_sequence(
        self,
        start_event: Union[str, Any],
        end_event: Union[str, Any],
        episode_lookup: Optional[Callable[[str], Any]] = None,
        detect_gaps: bool = True,
    ) -> TemporalSequence:
        """
        Reconstruct a sequence of events between two points.

        Args:
            start_event: Starting episode ID or Episode
            end_event: Ending episode ID or Episode
            episode_lookup: Function to look up episodes by ID
            detect_gaps: Whether to detect potential gaps in the sequence

        Returns:
            TemporalSequence object
        """
        # Get timestamps
        start_id = start_event if isinstance(start_event, str) else start_event.id
        end_id = end_event if isinstance(end_event, str) else end_event.id

        start_time = self.index.get_timestamp(start_id)
        end_time = self.index.get_timestamp(end_id)

        if start_time is None or end_time is None:
            return TemporalSequence(events=[], gap_count=0)

        # Ensure start < end
        if start_time > end_time:
            start_time, end_time = end_time, start_time
            start_id, end_id = end_id, start_id

        # Get all events in range
        results = self.index.query_range(start_time, end_time)
        events = self._to_temporal_events(results, start_time, episode_lookup)

        # Calculate duration
        duration = end_time - start_time

        # Detect gaps
        gap_count = 0
        if detect_gaps and len(events) > 1:
            for i in range(len(events) - 1):
                gap = events[i + 1].timestamp - events[i].timestamp
                if gap > self.max_gap_threshold:
                    gap_count += 1

        # Find start and end events
        start_ev = next((e for e in events if e.episode_id == start_id), None)
        end_ev = next((e for e in events if e.episode_id == end_id), None)

        return TemporalSequence(
            events=events,
            start_event=start_ev,
            end_event=end_ev,
            total_duration=duration,
            gap_count=gap_count,
        )

    def compute_temporal_decay_adjustment(
        self,
        episode: Any,
        base_lambda: float = 0.95,
        recency_window: timedelta = timedelta(days=7),
    ) -> float:
        """
        Compute adjusted decay rate based on recency.

        Recent memories decay slower.

        Args:
            episode: The episode
            base_lambda: Base decay rate
            recency_window: Window for considering an episode "recent"

        Returns:
            Adjusted lambda (higher = slower decay)
        """
        timestamp = getattr(episode, 'tocc', None)
        if timestamp is None:
            return base_lambda

        age = datetime.utcnow() - timestamp

        if age < recency_window:
            # Boost recent memories
            recency_factor = 1 - (age.total_seconds() / recency_window.total_seconds())
            boost = recency_factor * (self.recency_boost_factor - 1)
            return min(1.0, base_lambda * (1 + boost))

        return base_lambda

    def _normalize_reference(
        self,
        reference: Union[datetime, str, Any],
    ) -> Union[datetime, str]:
        """Normalize a reference to datetime or episode_id."""
        if isinstance(reference, datetime):
            return reference
        if isinstance(reference, str):
            # Try to parse as ISO datetime first
            try:
                return datetime.fromisoformat(reference)
            except ValueError:
                # Assume it's an episode ID
                return reference
        # Assume it's an Episode object
        return getattr(reference, 'id', reference)

    def _to_temporal_events(
        self,
        results: List[Tuple[datetime, str]],
        reference_time: Union[datetime, str],
        episode_lookup: Optional[Callable[[str], Any]] = None,
    ) -> List[TemporalEvent]:
        """Convert index results to TemporalEvent objects."""
        events = []

        # Get reference time for distance calculation
        if isinstance(reference_time, str):
            ref_time = self.index.get_timestamp(reference_time)
        else:
            ref_time = reference_time

        for timestamp, episode_id in results:
            content_preview = None
            trust = 1.0
            metadata = {}

            if episode_lookup:
                episode = episode_lookup(episode_id)
                if episode:
                    meta = getattr(episode, 'meta', {}) or {}
                    content_preview = (
                        meta.get('text') or
                        meta.get('content') or
                        meta.get('content_preview')
                    )
                    if content_preview and len(content_preview) > 200:
                        content_preview = content_preview[:200] + "..."
                    trust = getattr(episode, 'trust', 1.0)
                    metadata = meta

            temporal_distance = None
            if ref_time:
                temporal_distance = timestamp - ref_time

            events.append(TemporalEvent(
                episode_id=episode_id,
                timestamp=timestamp,
                content_preview=content_preview,
                trust=trust,
                metadata=metadata,
                temporal_distance=temporal_distance,
            ))

        return events

    def _rerank_by_similarity(
        self,
        events: List[TemporalEvent],
        query_vector: np.ndarray,
        episode_lookup: Callable[[str], Any],
        similarity_fn: Optional[Callable[[np.ndarray, np.ndarray], float]] = None,
    ) -> List[TemporalEvent]:
        """Re-rank events by semantic similarity to query."""
        scored_events = []

        for event in events:
            episode = episode_lookup(event.episode_id)
            if episode is None:
                scored_events.append((0.0, event))
                continue

            embedding = getattr(episode, 'e', None)
            if embedding is None:
                scored_events.append((0.0, event))
                continue

            if similarity_fn:
                sim = similarity_fn(query_vector, embedding)
            else:
                # Default cosine similarity
                norm_q = np.linalg.norm(query_vector)
                norm_e = np.linalg.norm(embedding)
                if norm_q == 0 or norm_e == 0:
                    sim = 0.0
                else:
                    sim = float(np.dot(query_vector, embedding) / (norm_q * norm_e))

            scored_events.append((sim, event))

        # Sort by similarity descending
        scored_events.sort(key=lambda x: x[0], reverse=True)

        return [event for _, event in scored_events]

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


def episodes_in_range(
    episodes: List[Any],
    start: Optional[datetime],
    end: Optional[datetime],
    limit: int = 50,
    ascending: bool = True,
) -> List[Any]:
    """
    Filter episodes by time range.

    Args:
        episodes: List of episodes
        start: Start time (inclusive)
        end: End time (inclusive)
        limit: Max results
        ascending: Sort order (true=oldest first, false=newest first)

    Returns:
        Filtered list of episodes
    """
    filtered = []
    for ep in episodes:
        ts = getattr(ep, 'tocc', None)
        if ts is None:
            continue

        if start and ts < start:
            continue
        if end and ts > end:
            continue

        filtered.append(ep)

    filtered.sort(key=lambda e: getattr(e, 'tocc', datetime.min), reverse=not ascending)
    return filtered[:limit]


def episodes_before(
    episodes: List[Any],
    timestamp: datetime,
    limit: int = 10,
) -> List[Any]:
    """
    Get episodes strictly before a timestamp.

    Args:
        episodes: List of episodes
        timestamp: Reference timestamp
        limit: Max results

    Returns:
        List of episodes, newest first
    """
    return episodes_in_range(
        episodes,
        start=None,
        end=timestamp - timedelta(microseconds=1),  # strict inequality
        limit=limit,
        ascending=False
    )


def episodes_after(
    episodes: List[Any],
    timestamp: datetime,
    limit: int = 10,
) -> List[Any]:
    """
    Get episodes strictly after a timestamp.

    Args:
        episodes: List of episodes
        timestamp: Reference timestamp
        limit: Max results

    Returns:
        List of episodes, oldest first
    """
    return episodes_in_range(
        episodes,
        start=timestamp + timedelta(microseconds=1),  # strict inequality
        end=None,
        limit=limit,
        ascending=True
    )
