"""
Tests for Cortex #513 — Temporal-Aware Retrieval Scoring.

Coverage:
  Unit tests:
    - parse_temporal_hints: ISO dates, relative terms, N-ago, in-month, unknowns
    - temporal_proximity_boost: exact match, half-life decay, bad inputs
    - filter_by_date: before/after bounds, undated items pass through
  Integration tests (mocked pipeline):
    - before_date / after_date filter out out-of-range claims
    - claims inside range survive, undated claims survive
    - results with date filters are sorted chronologically
    - temporal proximity boost raises score for near-date claims
    - temporal_boost_weight=0.0 produces no boost (non-temporal queries unaffected)
    - non-temporal query yields same results as before (no regression)
    - RetrievalConstraints carries before_date / after_date
"""

from __future__ import annotations

import asyncio
from dataclasses import dataclass, field
from datetime import datetime, timezone
from typing import Any, Dict, List, Optional
from unittest.mock import AsyncMock, MagicMock

import pytest

from cortex.core.retrieval import (
    RetrievalConstraints,
    RetrievalPipeline,
    _FTSResult,
    _VectorResult,
)
from cortex.core.temporal_scoring import (
    filter_by_date,
    parse_temporal_hints,
    temporal_proximity_boost,
)


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------

_FIXED_NOW = datetime(2024, 6, 15, 12, 0, 0, tzinfo=timezone.utc)


@dataclass
class MockConfig:
    retrieval_token_budget: int = 5000
    total_memory_token_budget: int = 10000
    cornerstone_token_budget: int = 800
    embedding_dim: int = 4
    embedding_model: str = "mock"
    retrieval_temporal_boost_weight: float = 0.15
    retrieval_rerank_enabled: bool = False
    retrieval_rerank_initial_candidates: int = 30


@dataclass
class MockClaim:
    """Minimal SemanticClaim stand-in carrying event_date."""
    claim_id: str
    subject: str = "Andrew"
    predicate: str = "likes"
    object_value: str = "coffee"
    event_date: Optional[str] = None
    status: str = "active"
    memory_class: str = "semantic"
    salience: str = "medium"
    cause: Optional[str] = None
    effect: Optional[str] = None
    state_before: Optional[str] = None
    state_after: Optional[str] = None


def make_storage(
    fts_items: Optional[List] = None,
    vector_items: Optional[List] = None,
    claims_by_id: Optional[Dict[str, MockClaim]] = None,
) -> MagicMock:
    storage = MagicMock()
    storage.embed = AsyncMock(return_value=[0.1, 0.2, 0.3, 0.4])

    _claims = [i for i in (fts_items or []) if getattr(i, "item_type", "claim") == "claim"]
    _events = [i for i in (fts_items or []) if getattr(i, "item_type", "claim") == "event"]

    async def _fts(query, table="claims", top_k=50, owner_id="default"):
        return _claims if table == "claims" else _events

    storage.fts_search = AsyncMock(side_effect=_fts)
    storage.vector_search = AsyncMock(return_value=vector_items or [])
    storage.log_retrieval = AsyncMock()

    _sens = {
        i.item_id: "normal"
        for i in (fts_items or []) + (vector_items or [])
        if getattr(i, "item_type", "claim") == "claim"
    }
    storage.get_claims_sensitivity_batch = AsyncMock(return_value=_sens)

    # get_claim: return from the claims_by_id dict if provided
    _cbi = claims_by_id or {}

    async def _get_claim(claim_id: str):
        return _cbi.get(claim_id)

    storage.get_claim = AsyncMock(side_effect=_get_claim)
    return storage


def make_token_budget() -> MagicMock:
    tb = MagicMock()
    tb.count_tokens = MagicMock(return_value=10)
    budget_obj = MagicMock()
    budget_obj.remaining = 5000
    tb.create_budget = MagicMock(return_value=budget_obj)
    return tb


def make_pipeline(storage, config=None, **kwargs) -> RetrievalPipeline:
    return RetrievalPipeline(
        storage,
        config or MockConfig(),
        make_token_budget(),
        **kwargs,
    )


# ---------------------------------------------------------------------------
# Unit tests: parse_temporal_hints
# ---------------------------------------------------------------------------

class TestParseTemporalHints:

    def test_iso_full_date(self):
        dt = parse_temporal_hints("what happened on 2024-03-15?", _now=_FIXED_NOW)
        assert dt == datetime(2024, 3, 15, tzinfo=timezone.utc)

    def test_iso_year_month(self):
        dt = parse_temporal_hints("events in 2024-03", _now=_FIXED_NOW)
        assert dt == datetime(2024, 3, 1, tzinfo=timezone.utc)

    def test_today(self):
        dt = parse_temporal_hints("what did I do today?", _now=_FIXED_NOW)
        assert dt == datetime(2024, 6, 15, tzinfo=timezone.utc)

    def test_yesterday(self):
        dt = parse_temporal_hints("yesterday's meeting", _now=_FIXED_NOW)
        assert dt == datetime(2024, 6, 14, tzinfo=timezone.utc)

    def test_last_week(self):
        dt = parse_temporal_hints("last week's tasks", _now=_FIXED_NOW)
        assert dt == datetime(2024, 6, 8, tzinfo=timezone.utc)

    def test_last_month(self):
        dt = parse_temporal_hints("last month's performance", _now=_FIXED_NOW)
        # June → May
        assert dt == datetime(2024, 5, 1, tzinfo=timezone.utc)

    def test_this_year(self):
        dt = parse_temporal_hints("this year's goals", _now=_FIXED_NOW)
        assert dt == datetime(2024, 1, 1, tzinfo=timezone.utc)

    def test_last_year(self):
        dt = parse_temporal_hints("last year revenue", _now=_FIXED_NOW)
        assert dt == datetime(2023, 1, 1, tzinfo=timezone.utc)

    def test_n_days_ago(self):
        dt = parse_temporal_hints("3 days ago sprint", _now=_FIXED_NOW)
        assert dt == datetime(2024, 6, 12, tzinfo=timezone.utc)

    def test_n_weeks_ago(self):
        dt = parse_temporal_hints("2 weeks ago decision", _now=_FIXED_NOW)
        assert dt == datetime(2024, 6, 1, tzinfo=timezone.utc)

    def test_in_month_name(self):
        dt = parse_temporal_hints("what happened in March?", _now=_FIXED_NOW)
        assert dt == datetime(2024, 3, 1, tzinfo=timezone.utc)

    def test_in_month_name_with_year(self):
        dt = parse_temporal_hints("events in March 2023", _now=_FIXED_NOW)
        assert dt == datetime(2023, 3, 1, tzinfo=timezone.utc)

    def test_no_hint_returns_none(self):
        dt = parse_temporal_hints("what is the capital of France?", _now=_FIXED_NOW)
        assert dt is None

    def test_empty_query(self):
        dt = parse_temporal_hints("", _now=_FIXED_NOW)
        assert dt is None

    def test_iso_date_takes_priority_over_relative(self):
        # Both ISO date and "last week" present — ISO wins
        dt = parse_temporal_hints("on 2024-01-01 last week's meeting", _now=_FIXED_NOW)
        assert dt == datetime(2024, 1, 1, tzinfo=timezone.utc)


# ---------------------------------------------------------------------------
# Unit tests: temporal_proximity_boost
# ---------------------------------------------------------------------------

class TestTemporalProximityBoost:

    def test_exact_match_is_one(self):
        ref = datetime(2024, 3, 15, tzinfo=timezone.utc)
        score = temporal_proximity_boost("2024-03-15", ref)
        assert abs(score - 1.0) < 1e-6

    def test_half_life_decay(self):
        ref = datetime(2024, 3, 15, tzinfo=timezone.utc)
        score = temporal_proximity_boost("2024-02-14", ref, half_life_days=30.0)
        # 29 days away ≈ 0.52 (half-life=30, boost = 2^(-29/30))
        assert 0.45 < score < 0.60

    def test_far_away_near_zero(self):
        ref = datetime(2024, 3, 15, tzinfo=timezone.utc)
        # 365 days away → boost = 2^(-365/30) ≈ very small
        score = temporal_proximity_boost("2023-03-15", ref, half_life_days=30.0)
        assert score < 0.01

    def test_future_date_also_decays(self):
        """Proximity is symmetric — future events also decay."""
        ref = datetime(2024, 3, 15, tzinfo=timezone.utc)
        score = temporal_proximity_boost("2024-04-14", ref, half_life_days=30.0)
        assert 0.45 < score < 0.60

    def test_year_month_format(self):
        """event_date may be YYYY-MM (treated as 1st of month)."""
        ref = datetime(2024, 3, 1, tzinfo=timezone.utc)
        score = temporal_proximity_boost("2024-03", ref)
        assert abs(score - 1.0) < 1e-6

    def test_bad_event_date_returns_zero(self):
        ref = datetime(2024, 3, 15, tzinfo=timezone.utc)
        assert temporal_proximity_boost("not-a-date", ref) == 0.0
        assert temporal_proximity_boost("", ref) == 0.0
        assert temporal_proximity_boost("2024", ref) == 0.0

    def test_boost_always_in_range(self):
        ref = datetime(2024, 3, 15, tzinfo=timezone.utc)
        for date_str in ["2024-03-15", "2023-01-01", "2025-12-31", "2024-03"]:
            b = temporal_proximity_boost(date_str, ref)
            assert 0.0 <= b <= 1.0, f"Out of range for {date_str}: {b}"


# ---------------------------------------------------------------------------
# Unit tests: filter_by_date
# ---------------------------------------------------------------------------

class TestFilterByDate:

    def test_no_filters_always_passes(self):
        assert filter_by_date("2024-03-15", None, None) is True
        assert filter_by_date(None, None, None) is True

    def test_undated_claim_always_passes(self):
        assert filter_by_date(None, "2024-01-01", "2024-12-31") is True
        assert filter_by_date("", "2024-01-01", "2024-12-31") is True

    def test_after_date_lower_bound(self):
        assert filter_by_date("2024-03-15", None, "2024-03-01") is True
        assert filter_by_date("2024-02-28", None, "2024-03-01") is False

    def test_before_date_upper_bound(self):
        assert filter_by_date("2024-03-15", "2024-04-01", None) is True
        assert filter_by_date("2024-04-15", "2024-04-01", None) is False

    def test_both_bounds_inside(self):
        assert filter_by_date("2024-03-15", "2024-04-01", "2024-03-01") is True

    def test_both_bounds_outside_after(self):
        assert filter_by_date("2024-02-01", "2024-04-01", "2024-03-01") is False

    def test_both_bounds_outside_before(self):
        assert filter_by_date("2024-05-01", "2024-04-01", "2024-03-01") is False

    def test_on_boundary_inclusive(self):
        # On the exact boundary — should pass (inclusive)
        assert filter_by_date("2024-03-01", "2024-04-01", "2024-03-01") is True
        assert filter_by_date("2024-04-01", "2024-04-01", "2024-03-01") is True

    def test_bad_date_strings_pass_through(self):
        # Unparseable filter dates → ignore filter → pass
        assert filter_by_date("2024-03-15", "not-a-date", None) is True
        assert filter_by_date("2024-03-15", None, "bad") is True


# ---------------------------------------------------------------------------
# Integration tests: pipeline with temporal scoring
# ---------------------------------------------------------------------------

class TestTemporalPipelineIntegration:
    """End-to-end tests with a mocked StorageAdapter."""

    def _make_fts_items(self, ids: List[str]) -> List[_FTSResult]:
        return [
            _FTSResult(item_id=iid, content=f"content {iid}", item_type="claim", rank=-1.0)
            for iid in ids
        ]

    def _make_claims_dict(self, specs: List[Dict]) -> Dict[str, MockClaim]:
        """Build a dict of MockClaim keyed by claim_id from spec dicts."""
        result = {}
        for spec in specs:
            cid = spec["id"]
            result[cid] = MockClaim(
                claim_id=cid,
                subject=spec.get("subject", "Andrew"),
                predicate=spec.get("predicate", "worked"),
                object_value=spec.get("object_value", "on project"),
                event_date=spec.get("event_date"),
            )
        return result

    @pytest.mark.asyncio
    async def test_date_filter_excludes_out_of_range_claims(self):
        """Claims outside after_date range are excluded."""
        item_ids = ["c1", "c2", "c3"]
        fts_items = self._make_fts_items(item_ids)
        claims = self._make_claims_dict([
            {"id": "c1", "event_date": "2024-01-15"},  # before March
            {"id": "c2", "event_date": "2024-03-15"},  # inside range
            {"id": "c3", "event_date": "2024-05-01"},  # inside range
        ])
        storage = make_storage(fts_items=fts_items, claims_by_id=claims)
        pipeline = make_pipeline(storage)

        constraints = RetrievalConstraints(after_date="2024-03-01")
        result = await pipeline.retrieve("meeting notes", constraints=constraints)

        ids = [i.item_id for i in result.items]
        assert "c1" not in ids, "c1 has event_date before after_date, should be excluded"
        assert "c2" in ids
        assert "c3" in ids

    @pytest.mark.asyncio
    async def test_before_date_filter_excludes_future_claims(self):
        """Claims after before_date are excluded."""
        item_ids = ["c1", "c2"]
        fts_items = self._make_fts_items(item_ids)
        claims = self._make_claims_dict([
            {"id": "c1", "event_date": "2024-01-15"},  # before cutoff → include
            {"id": "c2", "event_date": "2024-06-15"},  # after cutoff → exclude
        ])
        storage = make_storage(fts_items=fts_items, claims_by_id=claims)
        pipeline = make_pipeline(storage)

        constraints = RetrievalConstraints(before_date="2024-04-01")
        result = await pipeline.retrieve("old meeting", constraints=constraints)

        ids = [i.item_id for i in result.items]
        assert "c1" in ids
        assert "c2" not in ids

    @pytest.mark.asyncio
    async def test_undated_claims_survive_date_filters(self):
        """Claims without event_date are never excluded by date filters."""
        item_ids = ["dated", "undated"]
        fts_items = self._make_fts_items(item_ids)
        claims = self._make_claims_dict([
            {"id": "dated", "event_date": "2024-01-01"},   # before after_date → excluded
            {"id": "undated", "event_date": None},          # no date → always passes
        ])
        storage = make_storage(fts_items=fts_items, claims_by_id=claims)
        pipeline = make_pipeline(storage)

        constraints = RetrievalConstraints(after_date="2024-06-01")
        result = await pipeline.retrieve("info", constraints=constraints)

        ids = [i.item_id for i in result.items]
        assert "undated" in ids
        assert "dated" not in ids

    @pytest.mark.asyncio
    async def test_date_filters_sort_chronologically(self):
        """With date filters active, results are ordered by event_date ascending."""
        item_ids = ["c1", "c2", "c3"]
        fts_items = self._make_fts_items(item_ids)
        claims = self._make_claims_dict([
            {"id": "c1", "event_date": "2024-05-01"},
            {"id": "c2", "event_date": "2024-03-01"},
            {"id": "c3", "event_date": "2024-04-01"},
        ])
        storage = make_storage(fts_items=fts_items, claims_by_id=claims)
        pipeline = make_pipeline(storage)

        constraints = RetrievalConstraints(after_date="2024-01-01")
        result = await pipeline.retrieve("project timeline", constraints=constraints)

        ids = [i.item_id for i in result.items if i.item_type == "claim"]
        dates = [claims[i].event_date for i in ids]
        assert dates == sorted(dates), f"Expected chronological order, got {dates}"

    @pytest.mark.asyncio
    async def test_temporal_boost_raises_score_for_near_date(self):
        """Claims near the query's temporal hint get a higher score."""
        item_ids = ["near", "far"]
        fts_items = self._make_fts_items(item_ids)
        claims = self._make_claims_dict([
            {"id": "near", "event_date": "2024-03-15"},
            {"id": "far",  "event_date": "2023-01-01"},
        ])
        storage = make_storage(fts_items=fts_items, claims_by_id=claims)
        config = MockConfig(retrieval_temporal_boost_weight=0.15)
        pipeline = make_pipeline(storage, config=config)

        result = await pipeline.retrieve("events on 2024-03-15")
        scores = {i.item_id: i.score for i in result.items}

        if "near" in scores and "far" in scores:
            assert scores["near"] > scores["far"], (
                f"Near-date claim should score higher: near={scores['near']:.4f}, "
                f"far={scores['far']:.4f}"
            )

    @pytest.mark.asyncio
    async def test_zero_boost_weight_no_change(self):
        """With temporal_boost_weight=0.0, non-temporal queries are unaffected."""
        item_ids = ["c1", "c2"]
        fts_items = self._make_fts_items(item_ids)
        claims = self._make_claims_dict([
            {"id": "c1", "event_date": "2024-03-15"},
            {"id": "c2", "event_date": "2023-01-01"},
        ])
        storage = make_storage(fts_items=fts_items, claims_by_id=claims)
        config = MockConfig(retrieval_temporal_boost_weight=0.0)
        pipeline = make_pipeline(storage, config=config)

        result = await pipeline.retrieve("what did I work on?")
        # Just check no exceptions and items are returned
        assert len(result.items) >= 0  # result is valid

    @pytest.mark.asyncio
    async def test_non_temporal_query_no_regression(self):
        """Non-temporal query with no date filters works identically to before."""
        item_ids = ["c1", "c2"]
        fts_items = self._make_fts_items(item_ids)
        claims = self._make_claims_dict([
            {"id": "c1", "event_date": "2024-03-15"},
            {"id": "c2"},  # no event_date
        ])
        storage = make_storage(fts_items=fts_items, claims_by_id=claims)
        pipeline = make_pipeline(storage)

        result = await pipeline.retrieve("what is my favourite food?")
        ids = [i.item_id for i in result.items]
        # Both should appear — no filtering applied
        assert "c1" in ids
        assert "c2" in ids

    @pytest.mark.asyncio
    async def test_retrieval_constraints_carries_date_fields(self):
        """RetrievalConstraints dataclass correctly stores before/after dates."""
        c = RetrievalConstraints(before_date="2024-04-01", after_date="2024-03-01")
        assert c.before_date == "2024-04-01"
        assert c.after_date == "2024-03-01"

    @pytest.mark.asyncio
    async def test_no_date_filter_no_sort(self):
        """Without date filters, items are NOT sorted chronologically (score order)."""
        item_ids = ["c1", "c2", "c3"]
        # c3 comes first in fused (due to rank order)
        fts_items = [
            _FTSResult(item_id="c1", content="c1", item_type="claim", rank=-1.0),
            _FTSResult(item_id="c2", content="c2", item_type="claim", rank=-2.0),
            _FTSResult(item_id="c3", content="c3", item_type="claim", rank=-3.0),
        ]
        claims = self._make_claims_dict([
            {"id": "c1", "event_date": "2024-05-01"},
            {"id": "c2", "event_date": "2024-03-01"},
            {"id": "c3", "event_date": "2024-04-01"},
        ])
        storage = make_storage(fts_items=fts_items, claims_by_id=claims)
        config = MockConfig(retrieval_temporal_boost_weight=0.0)  # disable boost for clean test
        pipeline = make_pipeline(storage, config=config)

        result = await pipeline.retrieve("projects")
        # Without date filters, sorting is by relevance score, NOT date order
        # We can't assert exact order (depends on fusion) but at least no exception
        assert len(result.items) > 0

    @pytest.mark.asyncio
    async def test_both_date_filters_together(self):
        """Both before_date and after_date work together as a window."""
        item_ids = ["c1", "c2", "c3", "c4"]
        fts_items = self._make_fts_items(item_ids)
        claims = self._make_claims_dict([
            {"id": "c1", "event_date": "2024-01-01"},   # before window
            {"id": "c2", "event_date": "2024-03-15"},   # inside window
            {"id": "c3", "event_date": "2024-04-20"},   # inside window
            {"id": "c4", "event_date": "2024-06-01"},   # after window
        ])
        storage = make_storage(fts_items=fts_items, claims_by_id=claims)
        pipeline = make_pipeline(storage)

        constraints = RetrievalConstraints(
            after_date="2024-03-01",
            before_date="2024-05-01",
        )
        result = await pipeline.retrieve("Q1 work", constraints=constraints)

        ids = [i.item_id for i in result.items]
        assert "c1" not in ids, "c1 is before the window"
        assert "c2" in ids
        assert "c3" in ids
        assert "c4" not in ids, "c4 is after the window"
