"""
tests/test_retrieval_locomo_fields.py — Tests for LoCoMo v10 field surfacing.

Coverage:
  - _render_claim_content: baseline (no extra fields)
  - _render_claim_content: event_date prefix
  - _render_claim_content: status append (non-active only)
  - _render_claim_content: cause + effect (both, cause-only, effect-only)
  - _render_claim_content: state_before + state_after (both, single forms)
  - _render_claim_content: combined fields (event_date + cause + state)
  - RetrievedItem.metadata: memory_class, salience, status populated for claims
  - RetrievedItem.metadata: None for non-claim items
  - _hybrid_search: metadata attached from claim_cache (hydration path)
  - _hybrid_search: metadata attached for BM25 path (second-pass fetch)
"""

from __future__ import annotations

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

import pytest

from cortex.core.retrieval import (
    RetrievalConstraints,
    RetrievalPipeline,
    RetrievedItem,
    _render_claim_content,
)
from cortex.types import SemanticClaim


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

def make_claim(**kwargs) -> SemanticClaim:
    """Build a minimal SemanticClaim with overrideable fields."""
    defaults = dict(
        id="claim-test-001",
        owner_id="default",
        subject="Alice",
        predicate="likes",
        object_value="coffee",
        status="active",
        memory_class="semantic",
        salience="medium",
    )
    defaults.update(kwargs)
    return SemanticClaim(**defaults)


@dataclass
class MockConfig:
    retrieval_token_budget: int = 500
    total_memory_token_budget: int = 4000
    cornerstone_token_budget: int = 800
    embedding_dim: int = 4
    embedding_model: str = "mock"


@dataclass
class MockFTSResult:
    item_id: str
    content: str
    item_type: str = "claim"
    rank: float = -1.0


@dataclass
class MockVectorResult:
    item_id: str
    content: str = ""
    score: float = 0.9
    item_type: str = "claim"


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

    fts_items = fts_items or []
    vector_items = vector_items or []
    claims_by_id = claims_by_id or {}

    _claims = [i for i in fts_items if getattr(i, "item_type", "claim") == "claim"]
    _events = [i for i in fts_items 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)
    storage.log_retrieval = AsyncMock()

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

    async def _get_claim(claim_id, owner_id=None):
        return claims_by_id.get(claim_id)

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


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


def make_pipeline(storage=None, **kwargs) -> RetrievalPipeline:
    storage = storage or make_storage()
    config = kwargs.pop("config", MockConfig())
    token_budget = kwargs.pop("token_budget", make_token_budget())
    return RetrievalPipeline(storage, config, token_budget, **kwargs)


# ---------------------------------------------------------------------------
# Part A: _render_claim_content
# ---------------------------------------------------------------------------

class TestRenderClaimContent:
    """Tests for the _render_claim_content helper (pure function)."""

    def test_baseline_no_extra_fields(self):
        """No LoCoMo fields → plain 'subject predicate object'."""
        claim = make_claim()
        result = _render_claim_content(claim)
        assert result == "Alice likes coffee"

    def test_event_date_prefix(self):
        """event_date present → prefixed with [date]."""
        claim = make_claim(event_date="2024-03-15")
        result = _render_claim_content(claim)
        assert result.startswith("[2024-03-15] Alice likes coffee")

    def test_event_date_absent(self):
        """No event_date → no bracket prefix."""
        claim = make_claim(event_date=None)
        result = _render_claim_content(claim)
        assert not result.startswith("[")

    def test_status_active_not_appended(self):
        """status='active' → NOT appended."""
        claim = make_claim(status="active")
        result = _render_claim_content(claim)
        assert "[active]" not in result
        assert result == "Alice likes coffee"

    def test_status_superseded_appended(self):
        """status='superseded' → appended as [superseded]."""
        claim = make_claim(status="superseded")
        result = _render_claim_content(claim)
        assert "[superseded]" in result

    def test_status_completed_appended(self):
        """status='completed' → appended as [completed]."""
        claim = make_claim(status="completed")
        result = _render_claim_content(claim)
        assert "[completed]" in result

    def test_cause_and_effect_both(self):
        """Both cause and effect → '(because: cause → effect)'."""
        claim = make_claim(cause="stress", effect="insomnia")
        result = _render_claim_content(claim)
        assert "(because: stress \u2192 insomnia)" in result

    def test_cause_only(self):
        """Only cause → '(cause: cause)'."""
        claim = make_claim(cause="stress", effect=None)
        result = _render_claim_content(claim)
        assert "(cause: stress)" in result
        assert "effect" not in result
        assert "because" not in result

    def test_effect_only(self):
        """Only effect → '(effect: effect)'."""
        claim = make_claim(cause=None, effect="headache")
        result = _render_claim_content(claim)
        assert "(effect: headache)" in result
        assert "cause" not in result
        assert "because" not in result

    def test_cause_effect_absent(self):
        """No cause/effect → no because/cause/effect suffix."""
        claim = make_claim(cause=None, effect=None)
        result = _render_claim_content(claim)
        assert "because" not in result
        assert "cause" not in result
        assert "effect" not in result

    def test_state_before_and_after(self):
        """Both state_before and state_after → '(was: X → now: Y)'."""
        claim = make_claim(state_before="employed", state_after="freelance")
        result = _render_claim_content(claim)
        assert "(was: employed \u2192 now: freelance)" in result

    def test_state_before_only(self):
        """Only state_before → '(was: X)'."""
        claim = make_claim(state_before="employed", state_after=None)
        result = _render_claim_content(claim)
        assert "(was: employed)" in result
        assert "now:" not in result

    def test_state_after_only(self):
        """Only state_after → '(now: Y)'."""
        claim = make_claim(state_before=None, state_after="freelance")
        result = _render_claim_content(claim)
        assert "(now: freelance)" in result
        assert "was:" not in result

    def test_combined_all_fields(self):
        """All LoCoMo fields together → properly ordered output."""
        claim = make_claim(
            event_date="2024-06-01",
            status="completed",
            cause="promotion",
            effect="moved cities",
            state_before="London",
            state_after="Tokyo",
        )
        result = _render_claim_content(claim)
        assert result.startswith("[2024-06-01] Alice likes coffee")
        assert "[completed]" in result
        assert "(because: promotion \u2192 moved cities)" in result
        assert "(was: London \u2192 now: Tokyo)" in result

    def test_combined_partial(self):
        """event_date + cause-only → no state suffix."""
        claim = make_claim(
            event_date="2024-01-10",
            cause="boredom",
            state_before=None,
            state_after=None,
        )
        result = _render_claim_content(claim)
        assert result.startswith("[2024-01-10]")
        assert "(cause: boredom)" in result
        assert "was:" not in result
        assert "now:" not in result


# ---------------------------------------------------------------------------
# Part B: RetrievedItem.metadata
# ---------------------------------------------------------------------------

class TestRetrievedItemMetadata:
    """Tests for metadata field on RetrievedItem."""

    def test_metadata_field_exists(self):
        """RetrievedItem has a metadata field (None by default)."""
        item = RetrievedItem(
            item_type="claim",
            item_id="abc",
            content="some text",
            score=0.9,
        )
        assert hasattr(item, "metadata")
        assert item.metadata is None

    def test_metadata_populated_explicitly(self):
        """metadata can be set with a dict."""
        meta = {"memory_class": "episodic", "salience": "high", "status": "active"}
        item = RetrievedItem(
            item_type="claim",
            item_id="abc",
            content="some text",
            score=0.9,
            metadata=meta,
        )
        assert item.metadata["memory_class"] == "episodic"
        assert item.metadata["salience"] == "high"
        assert item.metadata["status"] == "active"


# ---------------------------------------------------------------------------
# Part B + Pipeline: metadata attached during _hybrid_search (hydration path)
# ---------------------------------------------------------------------------

class TestHybridSearchMetadata:
    """Integration-level tests verifying metadata is attached by _hybrid_search."""

    @pytest.mark.asyncio
    async def test_metadata_attached_via_hydration_path(self):
        """Vector-only result (empty content) → claim fetched → metadata attached."""
        claim = make_claim(
            id="c001",
            memory_class="episodic",
            salience="high",
            status="superseded",
        )
        storage = make_storage(
            fts_items=[],
            vector_items=[MockVectorResult(item_id="c001", content="", score=0.95)],
            claims_by_id={"c001": claim},
        )
        pipeline = make_pipeline(storage=storage)

        result = await pipeline.retrieve("test query", owner_id="default")
        assert result.items, "Expected at least one item"

        claim_item = next((i for i in result.items if i.item_id == "c001"), None)
        assert claim_item is not None, "Expected item c001 in results"
        assert claim_item.metadata is not None, "metadata should be populated"
        assert claim_item.metadata["memory_class"] == "episodic"
        assert claim_item.metadata["salience"] == "high"
        assert claim_item.metadata["status"] == "superseded"

    @pytest.mark.asyncio
    async def test_metadata_attached_via_bm25_path(self):
        """BM25 result (has content already) → second-pass fetch → metadata attached."""
        claim = make_claim(
            id="c002",
            memory_class="procedural",
            salience="low",
            status="active",
        )
        storage = make_storage(
            fts_items=[MockFTSResult(item_id="c002", content="Alice likes coffee")],
            vector_items=[],
            claims_by_id={"c002": claim},
        )
        pipeline = make_pipeline(storage=storage)

        result = await pipeline.retrieve("test query", owner_id="default")
        claim_item = next((i for i in result.items if i.item_id == "c002"), None)
        assert claim_item is not None, "Expected item c002 in results"
        assert claim_item.metadata is not None, "metadata should be populated"
        assert claim_item.metadata["memory_class"] == "procedural"
        assert claim_item.metadata["salience"] == "low"
        assert claim_item.metadata["status"] == "active"

    @pytest.mark.asyncio
    async def test_metadata_none_for_non_claim_items(self):
        """Non-claim items (e.g., event) → metadata is None."""
        storage = make_storage(
            fts_items=[MockFTSResult(item_id="e001", content="Some event", item_type="event")],
            vector_items=[],
        )
        # Override sensitivity batch so the event passes through
        storage.get_claims_sensitivity_batch = AsyncMock(return_value={})
        pipeline = make_pipeline(storage=storage)

        result = await pipeline.retrieve("test query", owner_id="default")
        event_item = next((i for i in result.items if i.item_id == "e001"), None)
        if event_item is not None:
            assert event_item.metadata is None

    @pytest.mark.asyncio
    async def test_hydration_renders_locomo_fields(self):
        """Hydrated claim content includes LoCoMo fields via _render_claim_content."""
        claim = make_claim(
            id="c003",
            event_date="2024-05-10",
            cause="deadline",
            effect="all-nighter",
            status="completed",
        )
        storage = make_storage(
            fts_items=[],
            vector_items=[MockVectorResult(item_id="c003", content="", score=0.8)],
            claims_by_id={"c003": claim},
        )
        pipeline = make_pipeline(storage=storage)

        result = await pipeline.retrieve("test query", owner_id="default")
        claim_item = next((i for i in result.items if i.item_id == "c003"), None)
        assert claim_item is not None
        assert "[2024-05-10]" in claim_item.content
        assert "[completed]" in claim_item.content
        assert "(because: deadline \u2192 all-nighter)" in claim_item.content

    @pytest.mark.asyncio
    async def test_metadata_keys_present(self):
        """Metadata dict always has memory_class, salience, and status keys."""
        claim = make_claim(id="c004")
        storage = make_storage(
            fts_items=[],
            vector_items=[MockVectorResult(item_id="c004", content="", score=0.7)],
            claims_by_id={"c004": claim},
        )
        pipeline = make_pipeline(storage=storage)

        result = await pipeline.retrieve("test query", owner_id="default")
        claim_item = next((i for i in result.items if i.item_id == "c004"), None)
        assert claim_item is not None
        assert claim_item.metadata is not None
        assert "memory_class" in claim_item.metadata
        assert "salience" in claim_item.metadata
        assert "status" in claim_item.metadata
