"""Tests for temporal reasoning."""

from dataclasses import dataclass, field
from datetime import datetime, timedelta
from typing import Any, Dict

from hieromem_core.reasoning.temporal import (
    TemporalEvent,
    TemporalSequence,
    TemporalIndex,
    TemporalReasoner,
)


@dataclass
class MockEpisode:
    """Mock episode for testing."""
    id: str
    tocc: datetime
    trust: float = 1.0
    meta: Dict[str, Any] = field(default_factory=dict)


class TestTemporalIndex:
    def test_add_and_get(self) -> None:
        index = TemporalIndex()
        ts = datetime(2025, 6, 15, 10, 0, 0)

        index.add("ep-1", ts)

        assert len(index) == 1
        assert index.get_timestamp("ep-1") == ts

    def test_remove(self) -> None:
        index = TemporalIndex()
        ts = datetime(2025, 6, 15)

        index.add("ep-1", ts)
        assert len(index) == 1

        result = index.remove("ep-1")
        assert result is True
        assert len(index) == 0
        assert index.get_timestamp("ep-1") is None

        # Remove nonexistent
        result = index.remove("ep-2")
        assert result is False

    def test_query_before(self) -> None:
        index = TemporalIndex()

        # Add episodes at different times
        index.add("ep-1", datetime(2025, 1, 1))
        index.add("ep-2", datetime(2025, 2, 1))
        index.add("ep-3", datetime(2025, 3, 1))
        index.add("ep-4", datetime(2025, 4, 1))

        # Query before ep-3
        results = index.query_before("ep-3", limit=10)
        ids = [ep_id for _, ep_id in results]

        assert "ep-1" in ids
        assert "ep-2" in ids
        assert "ep-3" not in ids
        assert "ep-4" not in ids

    def test_query_after(self) -> None:
        index = TemporalIndex()

        index.add("ep-1", datetime(2025, 1, 1))
        index.add("ep-2", datetime(2025, 2, 1))
        index.add("ep-3", datetime(2025, 3, 1))
        index.add("ep-4", datetime(2025, 4, 1))

        # Query after ep-2
        results = index.query_after("ep-2", limit=10)
        ids = [ep_id for _, ep_id in results]

        assert "ep-1" not in ids
        assert "ep-2" not in ids
        assert "ep-3" in ids
        assert "ep-4" in ids

    def test_query_range(self) -> None:
        index = TemporalIndex()

        index.add("ep-1", datetime(2025, 1, 1))
        index.add("ep-2", datetime(2025, 2, 1))
        index.add("ep-3", datetime(2025, 3, 1))
        index.add("ep-4", datetime(2025, 4, 1))

        # Query Feb-March
        results = index.query_range(
            datetime(2025, 2, 1),
            datetime(2025, 3, 31)
        )
        ids = [ep_id for _, ep_id in results]

        assert "ep-1" not in ids
        assert "ep-2" in ids
        assert "ep-3" in ids
        assert "ep-4" not in ids

    def test_query_around(self) -> None:
        index = TemporalIndex()

        index.add("ep-1", datetime(2025, 1, 1))
        index.add("ep-2", datetime(2025, 1, 15))
        index.add("ep-3", datetime(2025, 2, 1))
        index.add("ep-4", datetime(2025, 6, 1))

        # Query around ep-2 with 30 day window
        results = index.query_around("ep-2", window=timedelta(days=30))
        ids = [ep_id for _, ep_id in results]

        assert "ep-1" in ids
        assert "ep-2" in ids
        assert "ep-3" in ids
        assert "ep-4" not in ids


class TestTemporalReasoner:
    def test_index_episode(self) -> None:
        reasoner = TemporalReasoner()
        ep = MockEpisode("ep-1", datetime(2025, 6, 15))

        reasoner.index_episode(ep)

        assert len(reasoner.index) == 1

    def test_recall_before(self) -> None:
        reasoner = TemporalReasoner()

        episodes = [
            MockEpisode("ep-1", datetime(2025, 1, 1), meta={"text": "Event 1"}),
            MockEpisode("ep-2", datetime(2025, 2, 1), meta={"text": "Event 2"}),
            MockEpisode("ep-3", datetime(2025, 3, 1), meta={"text": "Event 3"}),
        ]

        for ep in episodes:
            reasoner.index_episode(ep)

        # Recall before ep-3
        events = reasoner.recall_before("ep-3", limit=10)

        assert len(events) == 2
        ids = [e.episode_id for e in events]
        assert "ep-1" in ids
        assert "ep-2" in ids

    def test_recall_after(self) -> None:
        reasoner = TemporalReasoner()

        episodes = [
            MockEpisode("ep-1", datetime(2025, 1, 1)),
            MockEpisode("ep-2", datetime(2025, 2, 1)),
            MockEpisode("ep-3", datetime(2025, 3, 1)),
        ]

        for ep in episodes:
            reasoner.index_episode(ep)

        events = reasoner.recall_after("ep-1", limit=10)

        assert len(events) == 2
        ids = [e.episode_id for e in events]
        assert "ep-2" in ids
        assert "ep-3" in ids

    def test_recall_temporal_range(self) -> None:
        reasoner = TemporalReasoner()

        episodes = [
            MockEpisode("ep-1", datetime(2025, 1, 1)),
            MockEpisode("ep-2", datetime(2025, 2, 1)),
            MockEpisode("ep-3", datetime(2025, 3, 1)),
            MockEpisode("ep-4", datetime(2025, 4, 1)),
        ]

        for ep in episodes:
            reasoner.index_episode(ep)

        events = reasoner.recall_temporal(
            start="2025-02-01",
            end="2025-03-31",
        )

        assert len(events) == 2
        ids = [e.episode_id for e in events]
        assert "ep-2" in ids
        assert "ep-3" in ids

    def test_reconstruct_sequence(self) -> None:
        reasoner = TemporalReasoner()

        episodes = [
            MockEpisode("ep-1", datetime(2025, 1, 1)),
            MockEpisode("ep-2", datetime(2025, 2, 1)),
            MockEpisode("ep-3", datetime(2025, 3, 1)),
        ]

        ep_map = {ep.id: ep for ep in episodes}
        for ep in episodes:
            reasoner.index_episode(ep)

        sequence = reasoner.reconstruct_sequence(
            start_event="ep-1",
            end_event="ep-3",
            episode_lookup=lambda x: ep_map.get(x),
        )

        assert len(sequence.events) == 3
        assert sequence.start_event is not None
        assert sequence.end_event is not None
        assert sequence.start_event.episode_id == "ep-1"
        assert sequence.end_event.episode_id == "ep-3"
        assert sequence.total_duration == timedelta(days=59)

    def test_gap_detection(self) -> None:
        reasoner = TemporalReasoner(max_gap_threshold=timedelta(days=30))

        episodes = [
            MockEpisode("ep-1", datetime(2025, 1, 1)),
            MockEpisode("ep-2", datetime(2025, 1, 15)),  # 14 days gap
            MockEpisode("ep-3", datetime(2025, 6, 1)),   # 137 days gap - should flag
        ]

        ep_map = {ep.id: ep for ep in episodes}
        for ep in episodes:
            reasoner.index_episode(ep)

        sequence = reasoner.reconstruct_sequence(
            start_event="ep-1",
            end_event="ep-3",
            episode_lookup=lambda x: ep_map.get(x),
        )

        assert sequence.gap_count == 1  # One gap > 30 days

    def test_temporal_decay_adjustment(self) -> None:
        reasoner = TemporalReasoner(recency_boost_factor=1.2)

        # Recent episode (1 day old)
        recent = MockEpisode("recent", datetime.utcnow() - timedelta(days=1))
        adj_recent = reasoner.compute_temporal_decay_adjustment(recent)
        assert adj_recent > 0.95  # Should be boosted

        # Old episode (30 days old)
        old = MockEpisode("old", datetime.utcnow() - timedelta(days=30))
        adj_old = reasoner.compute_temporal_decay_adjustment(old)
        assert adj_old == 0.95  # No boost


class TestTemporalEvent:
    def test_to_dict(self) -> None:
        event = TemporalEvent(
            episode_id="ep-1",
            timestamp=datetime(2025, 6, 15, 10, 30),
            content_preview="Test content",
            trust=0.9,
            temporal_distance=timedelta(hours=2),
        )

        d = event.to_dict()
        assert d["episode_id"] == "ep-1"
        assert d["trust"] == 0.9
        assert d["temporal_distance_seconds"] == 7200


class TestTemporalSequence:
    def test_to_dict(self) -> None:
        events = [
            TemporalEvent("ep-1", datetime(2025, 1, 1)),
            TemporalEvent("ep-2", datetime(2025, 2, 1)),
        ]

        seq = TemporalSequence(
            events=events,
            start_event=events[0],
            end_event=events[1],
            total_duration=timedelta(days=31),
        )

        d = seq.to_dict()
        assert d["event_count"] == 2
        assert d["total_duration_seconds"] == 31 * 24 * 3600
