"""
tests/test_access_tracking.py — Unit tests for access_count / last_accessed tracking.

Verifies that lightweight_retrieve and cosine_fallback fire-and-forget update
access_count and last_accessed on retrieved claims (#1447).

Coverage:
  1. access_count increments on retrieval via lightweight_retrieve
  2. last_accessed set to a recent UTC timestamp
  3. Non-claim items (episodes) are NOT tracked
  4. Tracking failure does not break retrieval (graceful degradation)
  5. cosine_fallback also fires access tracking for retrieved claims
"""

from __future__ import annotations

import asyncio
import struct
from datetime import datetime, timezone, timedelta
from typing import Any, Dict, List, Optional
from unittest.mock import AsyncMock, MagicMock, patch

import pytest

from cortex.core.retrieval_lightweight import lightweight_retrieve, _update_access_tracking
from cortex.core.retrieval_search import cosine_fallback
from cortex.core.retrieval_types import RetrievalConstraints


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

def _make_config(**overrides):
    cfg = MagicMock()
    cfg.retrieval_temporal_boost_weight = 0
    cfg.retrieval_token_budget = 3000
    cfg.retrieval_rerank_enabled = False
    cfg.retrieval_status_penalties = None
    cfg.owner_id = "default"
    for k, v in overrides.items():
        setattr(cfg, k, v)
    return cfg


def _make_constraints(**overrides):
    defaults = dict(
        top_k=10,
        sensitivity_max="high",
        memory_types=None,
        exclude_episodes=False,
        entity_ids=None,
        before_date=None,
        after_date=None,
        include_expired=True,
    )
    defaults.update(overrides)
    return RetrievalConstraints(**defaults)


def _pack_embedding(vec: List[float]) -> bytes:
    return struct.pack(f"{len(vec)}f", *vec)


# ---------------------------------------------------------------------------
# Minimal in-memory storage for access tracking tests
# ---------------------------------------------------------------------------

class _MockClaim:
    def __init__(self, claim_id, owner_id, content, access_count=0):
        self.id = claim_id
        self.owner_id = owner_id
        self.content = content
        self.access_count = access_count
        self.last_accessed = None
        self.status = "active"
        self.is_deleted = 0
        self.memory_class = "semantic"
        self.salience = "medium"
        self.explicitness = "explicit"
        self.stability = "stable"
        self.category = "test"
        self.created_at = "2026-01-01T00:00:00+00:00"
        self.updated_at = "2026-01-01T00:00:00+00:00"
        self.event_date = None
        self.valid_to = None
        self.metadata_json = None
        self.source_session = None
        self.subject = "test"
        self.predicate = "is"
        self.object_value = content
        self.confidence = 0.9


class AccessTrackingStorage:
    """In-memory storage for access tracking tests."""

    def __init__(self):
        self._claims: Dict[str, _MockClaim] = {}
        self._embeddings: Dict[str, List[float]] = {}
        self._update_calls: List[Dict] = []

    def add_claim(self, claim_id: str, owner_id: str, content: str,
                  embedding: Optional[List[float]] = None,
                  access_count: int = 0):
        self._claims[claim_id] = _MockClaim(claim_id, owner_id, content, access_count)
        if embedding:
            self._embeddings[claim_id] = embedding

    async def fts_search(self, query, table="claims", top_k=50, owner_id="default",
                         entity_id=None, entity_type=None, entity_ids=None):
        from cortex.types import FTSResult
        results = []
        for cid, claim in self._claims.items():
            if claim.owner_id != owner_id:
                continue
            if claim.is_deleted:
                continue
            results.append(FTSResult(item_id=cid, content=claim.content, rank=-1.0))
        return results[:top_k]

    async def vector_search(self, query_embedding, item_type=None, top_k=50,
                            owner_id="default"):
        from cortex.types import VectorResult
        results = []
        for cid, emb in self._embeddings.items():
            claim = self._claims.get(cid)
            if not claim or claim.owner_id != owner_id:
                continue
            sim = sum(a * b for a, b in zip(query_embedding, emb))
            results.append(VectorResult(item_type="claim", item_id=cid, score=sim))
        results.sort(key=lambda r: r.score, reverse=True)
        return results[:top_k]

    async def get_claim(self, claim_id: str) -> Optional[_MockClaim]:
        return self._claims.get(claim_id)

    async def update_claim(self, claim_id: str, owner_id: str = None, **kwargs) -> _MockClaim:
        claim = self._claims.get(claim_id)
        if not claim:
            raise ValueError(f"Claim {claim_id} not found")
        call_info = {"claim_id": claim_id, "owner_id": owner_id, **kwargs}
        self._update_calls.append(call_info)
        for k, v in kwargs.items():
            setattr(claim, k, v)
        return claim

    async def get_claims_sensitivity_batch(self, claim_ids):
        return {cid: "normal" for cid in claim_ids}

    async def log_retrieval(self, **kwargs):
        pass

    async def get_all_embeddings(self):
        rows = []
        for cid, emb in self._embeddings.items():
            blob = _pack_embedding(emb)
            rows.append({
                "item_type": "claim",
                "item_id": cid,
                "embedding": blob,
                "model_id": "test",
            })
        return rows


# ---------------------------------------------------------------------------
# Tests: _update_access_tracking (unit)
# ---------------------------------------------------------------------------

@pytest.mark.asyncio
async def test_update_access_tracking_increments_count():
    """_update_access_tracking increments access_count by 1."""
    storage = AccessTrackingStorage()
    storage.add_claim("claim-001", "owner-1", "Eva is the best AI", access_count=5)

    await _update_access_tracking(storage, ["claim-001"], "owner-1")

    updated_claim = storage._claims["claim-001"]
    assert updated_claim.access_count == 6


@pytest.mark.asyncio
async def test_update_access_tracking_sets_last_accessed():
    """_update_access_tracking sets last_accessed to a recent UTC timestamp."""
    storage = AccessTrackingStorage()
    storage.add_claim("claim-002", "owner-1", "Andrew builds cool things", access_count=0)

    before = datetime.now(timezone.utc)
    await _update_access_tracking(storage, ["claim-002"], "owner-1")
    after = datetime.now(timezone.utc)

    updated_claim = storage._claims["claim-002"]
    assert updated_claim.last_accessed is not None

    # Parse the stored timestamp
    ts_str = updated_claim.last_accessed
    # Handle both +00:00 and Z suffixes
    ts_str_normalized = ts_str.replace("+00:00", "").replace("Z", "")
    stored_ts = datetime.fromisoformat(ts_str_normalized).replace(tzinfo=timezone.utc)

    assert before - timedelta(seconds=1) <= stored_ts <= after + timedelta(seconds=1)


@pytest.mark.asyncio
async def test_update_access_tracking_batch():
    """_update_access_tracking handles multiple claim IDs."""
    storage = AccessTrackingStorage()
    storage.add_claim("claim-A", "owner-1", "Claim A content", access_count=0)
    storage.add_claim("claim-B", "owner-1", "Claim B content", access_count=3)
    storage.add_claim("claim-C", "owner-1", "Claim C content", access_count=10)

    await _update_access_tracking(storage, ["claim-A", "claim-B", "claim-C"], "owner-1")

    assert storage._claims["claim-A"].access_count == 1
    assert storage._claims["claim-B"].access_count == 4
    assert storage._claims["claim-C"].access_count == 11


@pytest.mark.asyncio
async def test_update_access_tracking_graceful_degradation():
    """_update_access_tracking failures do not raise — graceful degradation."""
    storage = AccessTrackingStorage()
    storage.add_claim("claim-good", "owner-1", "Good claim", access_count=0)

    # Make update_claim raise for the first call, succeed for the second
    original_update = storage.update_claim
    call_count = 0

    async def _flaky_update(claim_id, owner_id=None, **kwargs):
        nonlocal call_count
        call_count += 1
        if call_count == 1:
            raise RuntimeError("Simulated DB failure")
        return await original_update(claim_id, owner_id=owner_id, **kwargs)

    storage.update_claim = _flaky_update

    storage.add_claim("claim-bad", "owner-1", "Claim that fails update", access_count=0)

    # Should not raise even with a flaky storage
    await _update_access_tracking(storage, ["claim-bad", "claim-good"], "owner-1")

    # The second claim (claim-good) should still be updated
    assert storage._claims["claim-good"].access_count == 1


@pytest.mark.asyncio
async def test_update_access_tracking_missing_claim_is_skipped():
    """_update_access_tracking skips claims that don't exist in storage."""
    storage = AccessTrackingStorage()
    # Do NOT add any claims

    # Should complete without error even if claim doesn't exist
    await _update_access_tracking(storage, ["nonexistent-claim"], "owner-1")
    # No update calls for nonexistent claims
    assert storage._update_calls == []


# ---------------------------------------------------------------------------
# Tests: lightweight_retrieve fires access tracking
# ---------------------------------------------------------------------------

@pytest.mark.asyncio
async def test_lightweight_retrieve_fires_access_tracking():
    """lightweight_retrieve fires fire-and-forget access tracking for claims."""
    storage = AccessTrackingStorage()
    vec = [0.1, 0.2, 0.3, 0.4]
    storage.add_claim("claim-lw-1", "owner-1", "test memory content", embedding=vec,
                      access_count=0)

    config = _make_config()
    constraints = _make_constraints()

    async def _embed_fn(text):
        return vec

    items, _info = await lightweight_retrieve(
        storage=storage,
        config=config,
        query="test memory",
        constraints=constraints,
        owner_id="owner-1",
        rrf_k=60,
        candidate_multiplier=4,
        embed_fn=_embed_fn,
    )

    # Allow the background task to complete
    await asyncio.sleep(0.05)

    assert len(items) >= 1
    claim_ids = [it.item_id for it in items if it.item_type == "claim"]
    assert "claim-lw-1" in claim_ids

    # Verify access tracking was fired
    updated = storage._claims["claim-lw-1"]
    assert updated.access_count == 1
    assert updated.last_accessed is not None


@pytest.mark.asyncio
async def test_lightweight_retrieve_does_not_track_episodes():
    """lightweight_retrieve does NOT fire access tracking for non-claim items."""
    storage = AccessTrackingStorage()

    # Patch bm25_search to return an episode item
    from cortex.core.retrieval_types import _FTSResult

    with patch("cortex.core.retrieval_lightweight.bm25_search") as mock_bm25, \
         patch("cortex.core.retrieval_lightweight.vector_search") as mock_vec, \
         patch("cortex.core.retrieval_lightweight._update_access_tracking") as mock_tracker:

        mock_bm25.return_value = [
            _FTSResult(item_type="event", item_id="ep-001",
                       content="Episode content", rank=-1.0)
        ]
        mock_vec.return_value = []

        config = _make_config()
        constraints = _make_constraints()

        async def _embed_fn(text):
            return [0.1, 0.2, 0.3]

        items, _info = await lightweight_retrieve(
            storage=storage,
            config=config,
            query="episode query",
            constraints=constraints,
            owner_id="owner-1",
            rrf_k=60,
            candidate_multiplier=4,
            embed_fn=_embed_fn,
        )

        # _update_access_tracking should NOT be called for episode-only results
        # because _claim_ids_to_track would be empty
        if mock_tracker.called:
            # If called, verify no claim IDs were passed
            call_args = mock_tracker.call_args
            claim_ids_arg = call_args[0][1] if call_args[0] else call_args[1].get("claim_ids", [])
            assert claim_ids_arg == []


@pytest.mark.asyncio
async def test_lightweight_retrieve_tracking_failure_does_not_break_retrieval():
    """Tracking task failure does not affect the returned items."""
    from cortex.core.retrieval_types import _FTSResult

    storage = AccessTrackingStorage()
    storage.add_claim("claim-resilient", "owner-1", "resilient claim content", access_count=0)

    # Make update_claim always fail
    async def _failing_update(claim_id, owner_id=None, **kwargs):
        raise RuntimeError("Storage is down")

    storage.update_claim = _failing_update

    config = _make_config()
    constraints = _make_constraints()

    vec = [0.5, 0.5, 0.5]

    async def _embed_fn(text):
        return vec

    # Add embedding so vector search finds the claim
    storage._embeddings["claim-resilient"] = vec

    items, _info = await lightweight_retrieve(
        storage=storage,
        config=config,
        query="resilient claim",
        constraints=constraints,
        owner_id="owner-1",
        rrf_k=60,
        candidate_multiplier=4,
        embed_fn=_embed_fn,
    )

    # Allow background task to run and fail silently
    await asyncio.sleep(0.05)

    # Retrieval should still return results despite tracking failure
    assert isinstance(items, list)
    # The function returned without raising
    # (tracking failure is fire-and-forget — it must not propagate)


# ---------------------------------------------------------------------------
# Tests: cosine_fallback fires access tracking
# ---------------------------------------------------------------------------

@pytest.mark.asyncio
async def test_cosine_fallback_fires_access_tracking():
    """cosine_fallback fires fire-and-forget access tracking for retrieved claims."""
    storage = AccessTrackingStorage()
    vec = [1.0, 0.0, 0.0]
    storage.add_claim("claim-cf-1", "owner-1", "cosine fallback claim", embedding=vec,
                      access_count=0)

    query_embedding = [1.0, 0.0, 0.0]  # Perfect match

    results = await cosine_fallback(
        storage=storage,
        query_embedding=query_embedding,
        top_k=5,
        owner_id="owner-1",
    )

    # Allow background task to complete
    await asyncio.sleep(0.05)

    assert len(results) >= 1
    result_ids = [r.item_id for r in results]
    assert "claim-cf-1" in result_ids

    updated = storage._claims["claim-cf-1"]
    assert updated.access_count == 1
    assert updated.last_accessed is not None
