"""
tests/test_eval_harness.py — Unit tests for EvalHarness with mock plugin.

Strategy:
  - Mock the CortexPlugin entirely (no real DB, no real LLM)
  - Test the full harness pipeline: remember → retrieve → score → report
  - Test error handling: remember fails, retrieve fails
  - Test LLM judge integration
  - Test report structure and summary output
"""
from __future__ import annotations

import asyncio
import pytest
import pytest_asyncio
from dataclasses import dataclass, field
from typing import Any, List, Optional
from unittest.mock import AsyncMock, MagicMock

from cortex.eval.harness import EvalHarness, EvalReport, SessionResult, QuestionResult


# ---------------------------------------------------------------------------
# Mock plugin
# ---------------------------------------------------------------------------

@dataclass
class _MockRememberResult:
    session_id: str
    claims_stored: int = 2
    claims_extracted: int = 2
    ok: bool = True
    errors: List[str] = field(default_factory=list)


@dataclass
class _MockRetrievalResult:
    formatted_context: str = ""
    memories: List[Any] = field(default_factory=list)


class MockCortexPlugin:
    """A mock CortexPlugin that returns controllable responses."""

    def __init__(
        self,
        remember_result=None,
        retrieve_content: str = "",
        remember_raises: Optional[Exception] = None,
        retrieve_raises: Optional[Exception] = None,
    ):
        self._remember_result = remember_result or _MockRememberResult(session_id="test")
        self._retrieve_content = retrieve_content
        self._remember_raises = remember_raises
        self._retrieve_raises = retrieve_raises
        self.remember_calls = []
        self.retrieve_calls = []

    async def remember(self, conversation, session_id=None, scope="user", owner_id=None, **kwargs):
        self.remember_calls.append((conversation, session_id, scope))
        if self._remember_raises:
            raise self._remember_raises
        return self._remember_result

    async def retrieve(self, query, token_budget=None, constraints=None, mode="auto", owner_id=None, **kwargs):
        self.retrieve_calls.append((query, token_budget, constraints, mode))
        if self._retrieve_raises:
            raise self._retrieve_raises
        return _MockRetrievalResult(formatted_context=self._retrieve_content)


# ---------------------------------------------------------------------------
# LoCoMo harness tests
# ---------------------------------------------------------------------------

class TestEvalHarnessLocomo:
    """Tests for harness.run_locomo()."""

    @pytest.mark.asyncio
    async def test_run_locomo_sample(self):
        """Full pipeline with sample dataset, plugin returns empty answers."""
        plugin = MockCortexPlugin(retrieve_content="")
        harness = EvalHarness(plugin)
        report = await harness.run_locomo()

        assert isinstance(report, EvalReport)
        assert report.dataset_name == "locomo_sample"
        assert len(report.session_results) == 5
        assert report.metrics["total_questions"] == 20
        assert 0.0 <= report.metrics["exact_match_accuracy"] <= 1.0

    @pytest.mark.asyncio
    async def test_run_locomo_perfect_answers(self):
        """Plugin returns the exact expected answer for every question."""
        # We need to make the mock return the answer for each retrieve call
        # Since we can't know the question order, use a side_effect approach
        from cortex.eval.locomo import LoCoMoLoader

        loader = LoCoMoLoader()
        dataset = loader.load()

        # Build a lookup: question → answer
        qa_lookup = {}
        for session in dataset.sessions:
            for qa in session.qa_pairs:
                qa_lookup[qa.question] = qa.answer

        class SmartMockPlugin:
            async def remember(self, conversation, session_id=None, scope="user", owner_id=None, **kwargs):
                return _MockRememberResult(session_id=session_id or "test")

            async def retrieve(self, query, **kwargs):
                answer = qa_lookup.get(query, "")
                return _MockRetrievalResult(formatted_context=answer)

        harness = EvalHarness(SmartMockPlugin())
        report = await harness.run_locomo(dataset_dict={"sessions": [
            {
                "session_id": s.session_id,
                "conversation": s.conversation,
                "qa_pairs": [
                    {"question_id": qa.question_id, "question": qa.question, "answer": qa.answer, "type": qa.question_type}
                    for qa in s.qa_pairs
                ],
            }
            for s in dataset.sessions
        ]})

        # All answers are exact → exact_match_accuracy should be 1.0
        assert report.metrics["exact_match_accuracy"] == pytest.approx(1.0)
        assert report.metrics["fuzzy_match_accuracy"] == pytest.approx(1.0)
        assert report.metrics["avg_token_f1"] == pytest.approx(1.0)

    @pytest.mark.asyncio
    async def test_run_locomo_remember_fails(self):
        """remember() failure is recorded but evaluation continues."""
        plugin = MockCortexPlugin(
            remember_raises=RuntimeError("DB down"),
            retrieve_content="hiking",
        )
        harness = EvalHarness(plugin)
        report = await harness.run_locomo()

        # All sessions should still have results
        assert len(report.session_results) == 5
        for sr in report.session_results:
            assert sr.remember_error is not None
            assert "DB down" in sr.remember_error

    @pytest.mark.asyncio
    async def test_run_locomo_retrieve_fails(self):
        """retrieve() failure is recorded per question as empty answer."""
        plugin = MockCortexPlugin(retrieve_raises=RuntimeError("Retrieval down"))
        harness = EvalHarness(plugin)
        report = await harness.run_locomo()

        for sr in report.session_results:
            for qr in sr.question_results:
                assert qr.error is not None
                assert qr.predicted_answer == ""
                assert qr.exact_hit is False

    @pytest.mark.asyncio
    async def test_run_locomo_calls_remember_per_session(self):
        """remember() is called once per session."""
        plugin = MockCortexPlugin()
        harness = EvalHarness(plugin)
        await harness.run_locomo()

        assert len(plugin.remember_calls) == 5  # 5 sessions

    @pytest.mark.asyncio
    async def test_run_locomo_calls_retrieve_per_question(self):
        """retrieve() is called once per question."""
        plugin = MockCortexPlugin()
        harness = EvalHarness(plugin)
        await harness.run_locomo()

        assert len(plugin.retrieve_calls) == 20  # 20 questions total


# ---------------------------------------------------------------------------
# LongMemEval harness tests
# ---------------------------------------------------------------------------

class TestEvalHarnessLongMemEval:
    @pytest.mark.asyncio
    async def test_run_longmemeval_sample(self):
        plugin = MockCortexPlugin()
        harness = EvalHarness(plugin)
        report = await harness.run_longmemeval()

        assert isinstance(report, EvalReport)
        assert report.dataset_name == "longmemeval_sample"
        assert len(report.session_results) == 5
        assert report.metrics["total_questions"] > 0

    @pytest.mark.asyncio
    async def test_run_longmemeval_remember_called(self):
        plugin = MockCortexPlugin()
        harness = EvalHarness(plugin)
        await harness.run_longmemeval()
        assert len(plugin.remember_calls) == 5


# ---------------------------------------------------------------------------
# Answer extraction
# ---------------------------------------------------------------------------

class TestAnswerExtraction:
    def test_extract_formatted_context(self):
        harness = EvalHarness(MagicMock())
        result = MagicMock()
        result.formatted_context = "hiking"
        result.memories = []
        assert harness._extract_answer(result, "question") == "hiking"

    def test_extract_from_memories_list(self):
        harness = EvalHarness(MagicMock())

        @dataclass
        class MemObj:
            content: str

        result = MagicMock()
        result.formatted_context = ""
        result.memories = [MemObj(content="hiking"), MemObj(content="marathon")]
        extracted = harness._extract_answer(result, "question")
        assert "hiking" in extracted

    def test_extract_none_result(self):
        harness = EvalHarness(MagicMock())
        assert harness._extract_answer(None, "question") == ""

    def test_extract_fallback_to_str(self):
        harness = EvalHarness(MagicMock())
        result = MagicMock()
        result.formatted_context = ""
        result.memories = []
        result.context_window = "fallback answer"
        assert harness._extract_answer(result, "question") == "fallback answer"


# ---------------------------------------------------------------------------
# LLM judge integration
# ---------------------------------------------------------------------------

class TestLLMJudge:
    @pytest.mark.asyncio
    async def test_llm_judge_called(self):
        judge_calls = []

        def fake_judge(question, prediction, reference):
            judge_calls.append((question, prediction, reference))
            return 1.0 if prediction.lower() == reference.lower() else 0.0

        plugin = MockCortexPlugin(retrieve_content="hiking")
        harness = EvalHarness(plugin, llm_judge=fake_judge)
        report = await harness.run_locomo()

        assert len(judge_calls) == 20  # called for each question
        assert "avg_llm_judge_score" in report.metrics

    @pytest.mark.asyncio
    async def test_llm_judge_score_range(self):
        def fake_judge(question, prediction, reference):
            return 0.8

        plugin = MockCortexPlugin()
        harness = EvalHarness(plugin, llm_judge=fake_judge)
        report = await harness.run_locomo()

        score = report.metrics["avg_llm_judge_score"]
        assert 0.0 <= score <= 1.0


# ---------------------------------------------------------------------------
# EvalReport
# ---------------------------------------------------------------------------

class TestEvalReport:
    @pytest.mark.asyncio
    async def test_summary_output(self):
        plugin = MockCortexPlugin()
        harness = EvalHarness(plugin)
        report = await harness.run_locomo()
        summary = report.summary()
        assert "locomo_sample" in summary
        assert "Exact match" in summary
        assert "Fuzzy match" in summary

    @pytest.mark.asyncio
    async def test_to_dict_structure(self):
        plugin = MockCortexPlugin()
        harness = EvalHarness(plugin)
        report = await harness.run_locomo()
        d = report.to_dict()

        assert "dataset_name" in d
        assert "metrics" in d
        assert "sessions" in d
        assert "config" in d
        assert len(d["sessions"]) == 5

    @pytest.mark.asyncio
    async def test_metrics_keys(self):
        plugin = MockCortexPlugin()
        harness = EvalHarness(plugin)
        report = await harness.run_locomo()
        m = report.metrics

        for key in ["exact_match_accuracy", "fuzzy_match_accuracy", "avg_token_f1", "mrr",
                    "total_questions", "total_sessions", "retrieve_latency", "remember_latency"]:
            assert key in m, f"Missing metric: {key}"

    @pytest.mark.asyncio
    async def test_latency_recorded(self):
        plugin = MockCortexPlugin()
        harness = EvalHarness(plugin)
        report = await harness.run_locomo()

        lat = report.metrics["retrieve_latency"]
        assert lat["count"] == 20
        assert lat["mean"] >= 0.0
