"""AD-567b: Anchor-Aware Recall Formatting + Salience-Weighted Retrieval."""

from __future__ import annotations

import math
import time

import pytest

from probos.types import AnchorFrame, Episode, MemorySource, RecallScore


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

async def _start_episodic_memory(em: "EpisodicMemory") -> None:
    """Start EpisodicMemory, skipping if ChromaDB ONNX model is unavailable."""
    try:
        await em.start()
    except Exception as exc:
        if "INVALID_PROTOBUF" in str(exc) or "onnx" in str(exc).lower():
            pytest.skip(f"ChromaDB ONNX model unavailable: {exc}")
        raise


def _skip_on_onnx_error(func):
    """Decorator: skip test if ChromaDB ONNX model is corrupted/unavailable."""
    import functools

    @functools.wraps(func)
    async def wrapper(*args, **kwargs):
        try:
            return await func(*args, **kwargs)
        except Exception as exc:
            if "INVALID_PROTOBUF" in str(exc) or "onnx" in str(exc).lower():
                pytest.skip(f"ChromaDB ONNX model unavailable: {exc}")
            raise
    return wrapper

def _make_episode(
    *,
    user_input: str = "test input",
    timestamp: float | None = None,
    agent_ids: list[str] | None = None,
    source: str = "direct",
    anchors: AnchorFrame | None = None,
    reflection: str | None = None,
) -> Episode:
    return Episode(
        user_input=user_input,
        timestamp=timestamp or time.time(),
        agent_ids=agent_ids or ["agent-001"],
        source=source,
        anchors=anchors,
        reflection=reflection,
        outcomes=[{"intent": "test_intent", "success": True}],
    )


def _full_anchor() -> AnchorFrame:
    """AnchorFrame with all 10 fields filled."""
    return AnchorFrame(
        duty_cycle_id="duty-001",
        watch_section="alpha",
        channel="ward_room",
        channel_id="ch-123",
        department="science",
        participants=["Atlas", "Horizon"],
        trigger_agent="Atlas",
        trigger_type="ward_room_post",
        thread_id="thread-456",
        event_log_window=1000.0,
    )


def _half_anchor() -> AnchorFrame:
    """AnchorFrame with 5/10 fields filled."""
    return AnchorFrame(
        channel="ward_room",
        department="science",
        participants=["Atlas"],
        trigger_type="ward_room_post",
        trigger_agent="Atlas",
    )


# ---------------------------------------------------------------------------
# 1. RecallScore composite score formula
# ---------------------------------------------------------------------------

class TestRecallScoreComputation:
    def test_composite_score_formula(self):
        """Verify composite score matches AD-584c rebalanced formula with known inputs."""
        from probos.cognitive.episodic import EpisodicMemory

        ep = _make_episode(anchors=_full_anchor())
        rs = EpisodicMemory.score_recall(
            episode=ep,
            semantic_similarity=0.8,
            keyword_hits=3,
            trust_weight=0.9,
            hebbian_weight=0.7,
            recency_weight=0.6,
        )
        # AD-584c weights: 0.35*sem + 0.20*kw + 0.10*trust + 0.05*hebb + 0.15*rec + 0.15*anchor
        # + 0.10 convergence bonus (semantic > 0 AND keyword_hits > 0)
        # = 0.35*0.8 + 0.20*1.0 + 0.10*0.9 + 0.05*0.7 + 0.15*0.6 + 0.15*1.0 + 0.10
        expected = 0.35 * 0.8 + 0.20 * 1.0 + 0.10 * 0.9 + 0.05 * 0.7 + 0.15 * 0.6 + 0.15 * 1.0 + 0.10
        assert abs(rs.composite_score - expected) < 1e-9
        assert rs.anchor_confidence == 1.0

    def test_zero_inputs(self):
        """All-zero inputs produce composite_score = 0."""
        from probos.cognitive.episodic import EpisodicMemory

        ep = _make_episode()
        rs = EpisodicMemory.score_recall(
            episode=ep,
            semantic_similarity=0.0,
            keyword_hits=0,
            trust_weight=0.0,
            hebbian_weight=0.0,
            recency_weight=0.0,
        )
        assert rs.composite_score == 0.0


# ---------------------------------------------------------------------------
# 2. Recency weight exponential decay
# ---------------------------------------------------------------------------

class TestRecencyWeight:
    def test_fresh_episode(self):
        """Fresh episode (age ~0) ≈ 1.0 recency."""
        age_hours = 0.01
        rw = math.exp(-age_hours / 168.0)
        assert rw > 0.99

    def test_one_week_old(self):
        """~1-week-old episode ≈ 0.37 (1/e) recency."""
        age_hours = 168.0
        rw = math.exp(-age_hours / 168.0)
        assert abs(rw - 1 / math.e) < 0.01

    def test_two_weeks_old(self):
        """~2-week-old episode ≈ 0.14 recency."""
        age_hours = 336.0
        rw = math.exp(-age_hours / 168.0)
        assert abs(rw - math.exp(-2)) < 0.01


# ---------------------------------------------------------------------------
# 3. Anchor completeness scoring
# ---------------------------------------------------------------------------

class TestAnchorCompleteness:
    def test_all_filled(self):
        """10/10 fields = 1.0."""
        from probos.cognitive.episodic import EpisodicMemory

        ep = _make_episode(anchors=_full_anchor())
        rs = EpisodicMemory.score_recall(ep, semantic_similarity=0.5)
        assert rs.anchor_confidence == 1.0

    def test_half_filled(self):
        """Half-filled anchor → Johnson-weighted confidence ≈ 0.567."""
        from probos.cognitive.episodic import EpisodicMemory

        ep = _make_episode(anchors=_half_anchor())
        rs = EpisodicMemory.score_recall(ep, semantic_similarity=0.5)
        # _half_anchor: spatial=2/3, social=2/2, causal=1/1, temporal=0, evidential=0
        # = 0.25*(2/3) + 0.25*1.0 + 0.15*1.0 = 0.1667 + 0.25 + 0.15 = 0.5667
        expected = 0.25 * (2 / 3) + 0.25 * 1.0 + 0.15 * 1.0
        assert abs(rs.anchor_confidence - expected) < 1e-9

    def test_no_anchors(self):
        """anchors=None → 0.0."""
        from probos.cognitive.episodic import EpisodicMemory

        ep = _make_episode(anchors=None)
        rs = EpisodicMemory.score_recall(ep, semantic_similarity=0.5)
        assert rs.anchor_confidence == 0.0

    def test_empty_anchor(self):
        """All-default AnchorFrame → 0.0 (all fields empty/falsy)."""
        from probos.cognitive.episodic import EpisodicMemory

        ep = _make_episode(anchors=AnchorFrame())
        rs = EpisodicMemory.score_recall(ep, semantic_similarity=0.5)
        assert rs.anchor_confidence == 0.0


# ---------------------------------------------------------------------------
# 4. Budget enforcement
# ---------------------------------------------------------------------------

class TestBudgetEnforcement:
    @pytest.mark.asyncio
    @_skip_on_onnx_error
    async def test_budget_stops_accumulation(self, tmp_path):
        """recall_weighted() stops accumulating when context budget exceeded."""
        from probos.cognitive.episodic import EpisodicMemory

        em = EpisodicMemory(str(tmp_path / "ep.db"), max_episodes=100)
        await _start_episodic_memory(em)
        try:
            # Store 5 episodes with ~100 chars each
            for i in range(5):
                ep = _make_episode(
                    user_input=f"Episode content number {i} " + "x" * 80,
                    agent_ids=["agent-001"],
                )
                await em.store(ep)

            results = await em.recall_weighted(
                "agent-001", "Episode content",
                k=10, context_budget=250,
            )
            # Budget=250, each ~100 chars → should get ≤3
            assert len(results) <= 3
            # At least 1 always included
            assert len(results) >= 1
        finally:
            await em.stop()


# ---------------------------------------------------------------------------
# 5-8. FTS5 dual-write, eviction, seed, merge
# ---------------------------------------------------------------------------

class TestFTS5Integration:
    @pytest.mark.asyncio
    @_skip_on_onnx_error
    async def test_fts5_dual_write(self, tmp_path):
        """store() writes to FTS5; keyword_search finds it."""
        from probos.cognitive.episodic import EpisodicMemory

        em = EpisodicMemory(str(tmp_path / "ep.db"), max_episodes=100)
        await _start_episodic_memory(em)
        try:
            ep = _make_episode(user_input="quantum entanglement experiment results")
            await em.store(ep)

            results = await em.keyword_search("quantum entanglement", k=5)
            assert len(results) >= 1
            found_ids = [r[0] for r in results]
            assert ep.id in found_ids
        finally:
            await em.stop()

    @pytest.mark.asyncio
    @_skip_on_onnx_error
    async def test_fts5_eviction(self, tmp_path):
        """Evicted episodes are removed from FTS5."""
        from probos.cognitive.episodic import EpisodicMemory

        em = EpisodicMemory(str(tmp_path / "ep.db"), max_episodes=2)
        await _start_episodic_memory(em)
        try:
            ep1 = _make_episode(user_input="first experiment alpha", timestamp=100.0)
            ep2 = _make_episode(user_input="second experiment beta", timestamp=200.0)
            ep3 = _make_episode(user_input="third experiment gamma", timestamp=300.0)
            await em.store(ep1)
            await em.store(ep2)
            await em.store(ep3)  # Should evict ep1 (oldest)

            results = await em.keyword_search("alpha", k=5)
            found_ids = [r[0] for r in results]
            assert ep1.id not in found_ids
        finally:
            await em.stop()

    @pytest.mark.asyncio
    @_skip_on_onnx_error
    async def test_fts5_seed(self, tmp_path):
        """seed() populates FTS5; keyword_search finds seeded episodes."""
        from probos.cognitive.episodic import EpisodicMemory

        em = EpisodicMemory(str(tmp_path / "ep.db"), max_episodes=100)
        await _start_episodic_memory(em)
        try:
            ep = _make_episode(user_input="warp drive calibration anomaly")
            await em.seed([ep])

            # seed() swallows ONNX errors internally — skip if nothing was written
            if em._collection.count() == 0:
                pytest.skip("seed() silently failed due to ONNX model unavailability")

            results = await em.keyword_search("warp calibration", k=5)
            assert len(results) >= 1
            found_ids = [r[0] for r in results]
            assert ep.id in found_ids
        finally:
            await em.stop()

    @pytest.mark.asyncio
    @_skip_on_onnx_error
    async def test_fts5_semantic_merge(self, tmp_path):
        """Episode found by keyword but not vector still appears in recall_weighted()."""
        from probos.cognitive.episodic import EpisodicMemory

        em = EpisodicMemory(str(tmp_path / "ep.db"), max_episodes=100,
                            relevance_threshold=0.99)  # Very high threshold blocks semantic
        await _start_episodic_memory(em)
        try:
            # Store with very specific keyword content
            ep = _make_episode(
                user_input="xylophone zymurgy obscure keyword combination",
                agent_ids=["agent-001"],
            )
            await em.store(ep)

            # recall_weighted with keyword match query
            results = await em.recall_weighted(
                "agent-001", "xylophone zymurgy",
                k=5, context_budget=10000,
            )
            # May find it via FTS5 even if semantic threshold blocks it
            # (depends on ChromaDB embedding similarity — test validates the merge path)
            # At minimum, keyword_search should find it
            kw_results = await em.keyword_search("xylophone zymurgy", k=5)
            assert len(kw_results) >= 1
            # BF-134: If found via keyword-only merge, similarity should be fts_keyword_floor (0.2)
            # not 0.0 as it was before
            if results:
                for rs in results:
                    if rs.keyword_hits > 0 and rs.semantic_similarity == 0.0:
                        # This was the old behavior — should not happen now
                        pytest.fail("FTS-only hit should have floor similarity, not 0.0")
        finally:
            await em.stop()


# ---------------------------------------------------------------------------
# 9-10. Anchor-aware formatting
# ---------------------------------------------------------------------------

class TestAnchorAwareFormatting:
    def test_format_with_anchors(self):
        """_format_memory_section renders anchor header with channel, department, participants, trigger."""
        from probos.cognitive.cognitive_agent import CognitiveAgent

        agent = CognitiveAgent.__new__(CognitiveAgent)
        memories = [{
            "input": "Alert condition YELLOW detected",
            "source": "direct",
            "verified": True,
            "age": "2h 15m",
            "anchor_channel": "ward_room #security",
            "anchor_department": "Security",
            "anchor_participants": "Worf, Atlas",
            "anchor_trigger": "trust_variance",
        }]
        lines = agent._format_memory_section(memories)
        joined = "\n".join(lines)
        assert "[direct | verified]" in joined
        assert "2h 15m ago" in joined
        assert "ward_room #security" in joined
        assert "Security dept" in joined
        assert "with Worf, Atlas" in joined
        assert "re: trust_variance" in joined
        assert "Alert condition YELLOW detected" in joined

    def test_format_empty_anchors(self):
        """Graceful degradation when anchors are missing (old episodes)."""
        from probos.cognitive.cognitive_agent import CognitiveAgent

        agent = CognitiveAgent.__new__(CognitiveAgent)
        memories = [{
            "input": "Something happened",
            "source": "direct",
            "verified": False,
        }]
        lines = agent._format_memory_section(memories)
        joined = "\n".join(lines)
        assert "[direct | unverified]" in joined
        assert "Something happened" in joined
        # No anchor parts means no extra brackets
        assert "| |" not in joined


# ---------------------------------------------------------------------------
# 11-12. SECONDHAND source
# ---------------------------------------------------------------------------

class TestSecondhandSource:
    def test_secondhand_from_other_agent(self):
        """Episode from another agent's communication tagged SECONDHAND."""
        from probos.types import MemorySource

        # Validate the enum exists and has expected values
        assert MemorySource.SECONDHAND == "secondhand"
        assert MemorySource.DIRECT == "direct"

        ep = _make_episode(source=MemorySource.SECONDHAND)
        assert ep.source == "secondhand"

    def test_direct_preserved(self):
        """Episode from own action stays DIRECT."""
        ep = _make_episode(source=MemorySource.DIRECT)
        assert ep.source == "direct"


# ---------------------------------------------------------------------------
# 13. Recall ordering
# ---------------------------------------------------------------------------

class TestRecallOrdering:
    def test_higher_composite_ranks_first(self):
        """Higher composite score ranks first in sorted results."""
        from probos.cognitive.episodic import EpisodicMemory

        ep1 = _make_episode(user_input="low relevance")
        ep2 = _make_episode(user_input="high relevance", anchors=_full_anchor())

        rs1 = EpisodicMemory.score_recall(ep1, semantic_similarity=0.3, recency_weight=0.2)
        rs2 = EpisodicMemory.score_recall(ep2, semantic_similarity=0.9, recency_weight=0.9)

        results = sorted([rs1, rs2], key=lambda r: r.composite_score, reverse=True)
        assert results[0].episode.user_input == "high relevance"


# ---------------------------------------------------------------------------
# 14. Config weights
# ---------------------------------------------------------------------------

class TestConfigWeights:
    def test_custom_weights_affect_score(self):
        """Custom weights in MemoryConfig affect composite score."""
        from probos.cognitive.episodic import EpisodicMemory

        ep = _make_episode(anchors=_full_anchor())

        # Default weights
        rs_default = EpisodicMemory.score_recall(
            ep, semantic_similarity=0.8, trust_weight=0.9, recency_weight=0.5,
        )
        # Custom weights: boost trust, zero out semantic
        custom = {
            "semantic": 0.0, "keyword": 0.0, "trust": 0.80,
            "hebbian": 0.0, "recency": 0.10, "anchor": 0.10,
        }
        rs_custom = EpisodicMemory.score_recall(
            ep, semantic_similarity=0.8, trust_weight=0.9, recency_weight=0.5,
            weights=custom,
        )
        # With semantic zeroed, score should be lower (0.8 semantic was contributing)
        # and trust contribution should dominate
        assert rs_custom.composite_score != rs_default.composite_score
        # Trust dominates: 0.80 * 0.9 = 0.72 for trust alone
        expected_custom = 0.80 * 0.9 + 0.10 * 0.5 + 0.10 * 1.0
        assert abs(rs_custom.composite_score - expected_custom) < 1e-9


# ---------------------------------------------------------------------------
# 15. Backwards compatibility
# ---------------------------------------------------------------------------

class TestBackwardsCompatibility:
    @pytest.mark.asyncio
    @_skip_on_onnx_error
    async def test_recall_for_agent_still_returns_episodes(self, tmp_path):
        """recall_for_agent() still returns list[Episode] (not RecallScore)."""
        from probos.cognitive.episodic import EpisodicMemory

        em = EpisodicMemory(str(tmp_path / "ep.db"), max_episodes=100)
        await _start_episodic_memory(em)
        try:
            ep = _make_episode(user_input="backwards compatibility test")
            await em.store(ep)

            result = await em.recall_for_agent("agent-001", "backwards compatibility", k=5)
            assert isinstance(result, list)
            if result:
                assert isinstance(result[0], Episode)
        finally:
            await em.stop()

    @pytest.mark.asyncio
    async def test_recent_for_agent_unchanged(self, tmp_path):
        """recent_for_agent() still works as before."""
        from probos.cognitive.episodic import EpisodicMemory

        em = EpisodicMemory(str(tmp_path / "ep.db"), max_episodes=100)
        await _start_episodic_memory(em)
        try:
            ep = _make_episode(user_input="recent test episode")
            await em.store(ep)

            result = await em.recent_for_agent("agent-001", k=5)
            assert isinstance(result, list)
            if result:
                assert isinstance(result[0], Episode)
        finally:
            await em.stop()


# ---------------------------------------------------------------------------
# RecallScore dataclass validation
# ---------------------------------------------------------------------------

class TestRecallScoreDataclass:
    def test_frozen(self):
        """RecallScore is immutable."""
        ep = _make_episode()
        rs = RecallScore(episode=ep, composite_score=0.5)
        with pytest.raises(AttributeError):
            rs.composite_score = 0.9  # type: ignore[misc]

    def test_defaults(self):
        """RecallScore defaults are sensible."""
        ep = _make_episode()
        rs = RecallScore(episode=ep)
        assert rs.semantic_similarity == 0.0
        assert rs.keyword_hits == 0
        assert rs.trust_weight == 0.5
        assert rs.hebbian_weight == 0.5
        assert rs.recency_weight == 0.0
        assert rs.anchor_confidence == 0.0
        assert rs.composite_score == 0.0


# ---------------------------------------------------------------------------
# BF-134: Configurable Recall Threshold + FTS Floor
# ---------------------------------------------------------------------------

class TestBF134ThresholdAndFloor:
    """BF-134: Verify configurable agent recall threshold and FTS keyword floor."""

    def test_agent_recall_threshold_default(self):
        """Default agent_recall_threshold is 0.25 (AD-593 raised from 0.15)."""
        from probos.cognitive.episodic import EpisodicMemory
        em = EpisodicMemory("/dev/null", max_episodes=100)
        assert em._agent_recall_threshold == 0.25

    def test_agent_recall_threshold_configurable(self):
        """Custom agent_recall_threshold is stored and used."""
        from probos.cognitive.episodic import EpisodicMemory
        em = EpisodicMemory("/dev/null", max_episodes=100, agent_recall_threshold=0.1)
        assert em._agent_recall_threshold == 0.1

    def test_fts_keyword_floor_default(self):
        """Default fts_keyword_floor is 0.2."""
        from probos.cognitive.episodic import EpisodicMemory
        em = EpisodicMemory("/dev/null", max_episodes=100)
        assert em._fts_keyword_floor == 0.2

    def test_fts_keyword_floor_configurable(self):
        """Custom fts_keyword_floor is stored."""
        from probos.cognitive.episodic import EpisodicMemory
        em = EpisodicMemory("/dev/null", max_episodes=100, fts_keyword_floor=0.3)
        assert em._fts_keyword_floor == 0.3

    @pytest.mark.asyncio
    @_skip_on_onnx_error
    async def test_fts_keyword_floor_boosts_composite_score(self, tmp_path):
        """FTS-only hits get fts_keyword_floor as similarity, contributing to composite score."""
        from probos.cognitive.episodic import EpisodicMemory

        em = EpisodicMemory(
            str(tmp_path / "fts_floor.db"), max_episodes=100,
            relevance_threshold=0.99,  # Block semantic results
            fts_keyword_floor=0.2,
        )
        await _start_episodic_memory(em)
        try:
            ep = _make_episode(
                user_input="xylophone zymurgy unique keyword test",
                agent_ids=["agent-001"],
            )
            await em.store(ep)
            results = await em.recall_weighted(
                "agent-001", "xylophone zymurgy",
                k=5, context_budget=10000,
            )
            # Find via keyword
            kw_results = await em.keyword_search("xylophone zymurgy", k=5)
            assert len(kw_results) >= 1
            # If recall_weighted found it, the semantic component should be 0.2 not 0.0
            if results:
                for rs in results:
                    if rs.keyword_hits > 0:
                        # semantic weight is 0.35, so contribution = 0.35 * 0.2 = 0.07
                        assert rs.composite_score > 0.0
        finally:
            await em.stop()

    def test_threshold_and_anchor_gate_work_together(self):
        """Both agent_recall_threshold and anchor_confidence_gate are independent filters."""
        from probos.config import MemoryConfig
        cfg = MemoryConfig()
        # Verify they are independent config fields with correct defaults
        assert cfg.agent_recall_threshold == 0.25
        assert cfg.anchor_confidence_gate == 0.3
        assert cfg.fts_keyword_semantic_floor == 0.2
        # Threshold is for semantic similarity, gate is for anchor confidence
        # Both must pass for an episode to be returned
