"""
cortex/eval/harness.py — EvalHarness: Benchmark evaluation driver.

Orchestrates the full evaluation pipeline:
  1. Load benchmark dataset (LoCoMo or LongMemEval format)
  2. Run remember() on each session's conversation
  3. Run retrieve() for each question
  4. Score answers (exact match, fuzzy match, optional LLM-as-judge)
  5. Report metrics: precision, recall, F1, MRR, latency

Usage::

    from cortex.api.plugin import CortexPlugin
    from cortex.eval.harness import EvalHarness

    plugin = CortexPlugin.from_config("cortex.toml")
    harness = EvalHarness(plugin)
    results = await harness.run_locomo()
    print(results.summary())
"""
from __future__ import annotations

import asyncio
import logging
import time
from dataclasses import dataclass, field
from typing import Any, Callable, Dict, List, Optional

from .metrics import (
    compute_qa_metrics,
    latency_stats,
    mean_reciprocal_rank,
    precision_recall_f1,
    token_f1_score,
    fuzzy_match,
    exact_match,
)

logger = logging.getLogger(__name__)


# ---------------------------------------------------------------------------
# Result types
# ---------------------------------------------------------------------------

@dataclass
class QuestionResult:
    """Result for a single QA evaluation."""
    question_id: str
    question: str
    expected_answer: str
    predicted_answer: str
    exact_hit: bool
    fuzzy_hit: bool
    token_f1: float
    llm_judge_score: Optional[float]
    remember_latency_ms: float
    retrieve_latency_ms: float
    error: Optional[str] = None


@dataclass
class SessionResult:
    """Result for a single session evaluation."""
    session_id: str
    question_results: List[QuestionResult] = field(default_factory=list)
    remember_latency_ms: float = 0.0
    claims_stored: int = 0
    remember_error: Optional[str] = None

    @property
    def n_questions(self) -> int:
        return len(self.question_results)

    @property
    def exact_match_count(self) -> int:
        return sum(1 for r in self.question_results if r.exact_hit)

    @property
    def fuzzy_match_count(self) -> int:
        return sum(1 for r in self.question_results if r.fuzzy_hit)


@dataclass
class EvalReport:
    """
    Full evaluation report with aggregate metrics.

    Attributes:
        dataset_name: Name of the benchmark dataset.
        session_results: Per-session results.
        metrics: Aggregate metrics dict.
        config: Eval configuration used.
    """
    dataset_name: str
    session_results: List[SessionResult] = field(default_factory=list)
    metrics: Dict[str, Any] = field(default_factory=dict)
    config: Dict[str, Any] = field(default_factory=dict)

    def summary(self, indent: int = 0) -> str:
        """Return a human-readable summary of the eval results."""
        pad = " " * indent
        lines = [
            f"{pad}=== Eval Report: {self.dataset_name} ===",
            f"{pad}Sessions evaluated: {len(self.session_results)}",
            f"{pad}Total questions:    {self.metrics.get('total_questions', 0)}",
            f"{pad}Exact match:        {self.metrics.get('exact_match_accuracy', 0):.1%}",
            f"{pad}Fuzzy match:        {self.metrics.get('fuzzy_match_accuracy', 0):.1%}",
            f"{pad}Avg token F1:       {self.metrics.get('avg_token_f1', 0):.3f}",
            f"{pad}MRR:                {self.metrics.get('mrr', 0):.3f}",
        ]
        if "avg_llm_judge_score" in self.metrics:
            lines.append(f"{pad}LLM judge score:    {self.metrics['avg_llm_judge_score']:.3f}")
        if "retrieve_latency" in self.metrics:
            lat = self.metrics["retrieve_latency"]
            lines.append(f"{pad}Retrieve p95 (ms):  {lat.get('p95', 0):.1f}")
        if "remember_latency" in self.metrics:
            lat = self.metrics["remember_latency"]
            lines.append(f"{pad}Remember p95 (ms):  {lat.get('p95', 0):.1f}")
        return "\n".join(lines)

    def to_dict(self) -> Dict[str, Any]:
        """Return the report as a plain dict (JSON-serialisable)."""
        return {
            "dataset_name": self.dataset_name,
            "metrics": self.metrics,
            "sessions": [
                {
                    "session_id": sr.session_id,
                    "n_questions": sr.n_questions,
                    "exact_match_count": sr.exact_match_count,
                    "fuzzy_match_count": sr.fuzzy_match_count,
                    "claims_stored": sr.claims_stored,
                    "remember_latency_ms": sr.remember_latency_ms,
                    "remember_error": sr.remember_error,
                    "questions": [
                        {
                            "question_id": r.question_id,
                            "question": r.question,
                            "expected": r.expected_answer,
                            "predicted": r.predicted_answer,
                            "exact_hit": r.exact_hit,
                            "fuzzy_hit": r.fuzzy_hit,
                            "token_f1": r.token_f1,
                            "llm_judge_score": r.llm_judge_score,
                            "retrieve_latency_ms": r.retrieve_latency_ms,
                            "error": r.error,
                        }
                        for r in sr.question_results
                    ],
                }
                for sr in self.session_results
            ],
            "config": self.config,
        }


# ---------------------------------------------------------------------------
# EvalHarness
# ---------------------------------------------------------------------------

class EvalHarness:
    """
    Drives the full benchmark evaluation against a CortexPlugin instance.

    Args:
        plugin:         A CortexPlugin (or duck-typed compatible) instance.
        fuzzy_threshold: Threshold for fuzzy match scoring (default 0.6).
        llm_judge:      Optional callable(question, prediction, reference) -> float.
                        If provided, adds LLM-as-judge scores to results.
        max_retrieved:  Max retrieved items to inspect per question (for MRR).
    """

    def __init__(
        self,
        plugin: Any,
        fuzzy_threshold: float = 0.6,
        llm_judge: Optional[Callable[[str, str, str], float]] = None,
        max_retrieved: int = 10,
    ):
        self.plugin = plugin
        self.fuzzy_threshold = fuzzy_threshold
        self.llm_judge = llm_judge
        self.max_retrieved = max_retrieved

    # ------------------------------------------------------------------
    # Public runners
    # ------------------------------------------------------------------

    async def run_locomo(
        self,
        path: Optional[str] = None,
        dataset_dict: Optional[Dict[str, Any]] = None,
    ) -> EvalReport:
        """
        Run evaluation on a LoCoMo dataset.

        Args:
            path: Path to LoCoMo JSON file. If None uses bundled sample.
            dataset_dict: Pre-parsed dataset dict (alternative to path).

        Returns:
            EvalReport
        """
        from .locomo import LoCoMoLoader, LoCoMoDataset

        loader = LoCoMoLoader()
        if dataset_dict is not None:
            dataset = loader.load_dict(dataset_dict)
        else:
            dataset = loader.load(path)

        logger.info("Running LoCoMo eval on %d sessions", len(dataset.sessions))
        session_results = []

        for session in dataset.sessions:
            sr = await self._eval_locomo_session(session)
            session_results.append(sr)

        return self._build_report(dataset.name, session_results)

    async def run_longmemeval(
        self,
        path: Optional[str] = None,
        dataset_dict: Optional[Dict[str, Any]] = None,
    ) -> EvalReport:
        """
        Run evaluation on a LongMemEval dataset.

        Args:
            path: Path to LongMemEval JSON file. If None uses bundled sample.
            dataset_dict: Pre-parsed dataset dict (alternative to path).

        Returns:
            EvalReport
        """
        from .longmemeval import LongMemEvalLoader, LongMemEvalDataset

        loader = LongMemEvalLoader()
        if dataset_dict is not None:
            dataset = loader.load_dict(dataset_dict)
        else:
            dataset = loader.load(path)

        logger.info("Running LongMemEval on %d sessions", len(dataset.sessions))
        session_results = []

        for session in dataset.sessions:
            sr = await self._eval_longmemeval_session(session)
            session_results.append(sr)

        return self._build_report(dataset.name, session_results)

    # ------------------------------------------------------------------
    # Session-level evaluation
    # ------------------------------------------------------------------

    async def _eval_locomo_session(self, session: Any) -> SessionResult:
        """Evaluate a single LoCoMo session."""
        sr = SessionResult(session_id=session.session_id)

        # 1. Remember the conversation
        t0 = time.perf_counter()
        try:
            remember_result = await self.plugin.remember(
                conversation=session.messages,
                session_id=session.session_id,
            )
            sr.remember_latency_ms = (time.perf_counter() - t0) * 1000
            sr.claims_stored = getattr(remember_result, "claims_stored", 0)
        except Exception as exc:
            sr.remember_latency_ms = (time.perf_counter() - t0) * 1000
            sr.remember_error = str(exc)
            logger.warning("remember() failed for session %s: %s", session.session_id, exc)

        # 2. Evaluate each QA pair
        for qa in session.qa_pairs:
            qr = await self._eval_question(
                question_id=qa.question_id,
                question=qa.question,
                expected_answer=qa.answer,
                remember_latency_ms=sr.remember_latency_ms,
            )
            sr.question_results.append(qr)

        return sr

    async def _eval_longmemeval_session(self, session: Any) -> SessionResult:
        """Evaluate a single LongMemEval session."""
        sr = SessionResult(session_id=session.session_id)

        # 1. Remember the conversation
        t0 = time.perf_counter()
        try:
            remember_result = await self.plugin.remember(
                conversation=session.conversation,
                session_id=session.session_id,
            )
            sr.remember_latency_ms = (time.perf_counter() - t0) * 1000
            sr.claims_stored = getattr(remember_result, "claims_stored", 0)
        except Exception as exc:
            sr.remember_latency_ms = (time.perf_counter() - t0) * 1000
            sr.remember_error = str(exc)
            logger.warning("remember() failed for session %s: %s", session.session_id, exc)

        # 2. Evaluate each question
        for q in session.questions:
            qr = await self._eval_question(
                question_id=q.question_id,
                question=q.question,
                expected_answer=q.answer,
                remember_latency_ms=sr.remember_latency_ms,
            )
            sr.question_results.append(qr)

        return sr

    async def _eval_question(
        self,
        question_id: str,
        question: str,
        expected_answer: str,
        remember_latency_ms: float,
    ) -> QuestionResult:
        """Run retrieve() for a single question and score the result."""
        predicted_answer = ""
        retrieve_latency_ms = 0.0
        error = None

        t0 = time.perf_counter()
        try:
            result = await self.plugin.retrieve(query=question)
            retrieve_latency_ms = (time.perf_counter() - t0) * 1000
            predicted_answer = self._extract_answer(result, question)
        except Exception as exc:
            retrieve_latency_ms = (time.perf_counter() - t0) * 1000
            error = str(exc)
            logger.warning("retrieve() failed for question %s: %s", question_id, exc)

        # Score
        exact_hit = exact_match(predicted_answer, expected_answer)
        fuzzy_hit = fuzzy_match(predicted_answer, expected_answer, self.fuzzy_threshold)
        f1 = token_f1_score(predicted_answer, expected_answer)

        llm_score = None
        if self.llm_judge is not None:
            try:
                llm_score = self.llm_judge(question, predicted_answer, expected_answer)
            except Exception as exc:
                logger.warning("LLM judge failed for %s: %s", question_id, exc)

        return QuestionResult(
            question_id=question_id,
            question=question,
            expected_answer=expected_answer,
            predicted_answer=predicted_answer,
            exact_hit=exact_hit,
            fuzzy_hit=fuzzy_hit,
            token_f1=f1,
            llm_judge_score=llm_score,
            remember_latency_ms=remember_latency_ms,
            retrieve_latency_ms=retrieve_latency_ms,
            error=error,
        )

    # ------------------------------------------------------------------
    # Answer extraction from RetrievalResult
    # ------------------------------------------------------------------

    def _extract_answer(self, result: Any, question: str) -> str:
        """
        Extract a text answer from the plugin's RetrievalResult.

        Tries common attribute patterns; falls back to str().
        """
        if result is None:
            return ""

        # RetrievalResult.formatted_context — the injected memory text
        if hasattr(result, "formatted_context") and result.formatted_context:
            return result.formatted_context

        # RetrievalResult.memories — list of memory objects
        if hasattr(result, "memories") and result.memories:
            parts = []
            for mem in result.memories[: self.max_retrieved]:
                if hasattr(mem, "content"):
                    parts.append(str(mem.content))
                elif isinstance(mem, dict):
                    parts.append(str(mem.get("content", mem)))
                else:
                    parts.append(str(mem))
            return " ".join(parts)

        # RetrievalResult.context_window / text / answer
        for attr in ("context_window", "text", "answer", "content"):
            val = getattr(result, attr, None)
            if val:
                return str(val)

        return str(result) if result else ""

    # ------------------------------------------------------------------
    # Report builder
    # ------------------------------------------------------------------

    def _build_report(
        self,
        dataset_name: str,
        session_results: List[SessionResult],
    ) -> EvalReport:
        """Aggregate session results into an EvalReport."""
        all_question_results = [
            qr
            for sr in session_results
            for qr in sr.question_results
        ]

        n = len(all_question_results)
        predictions = [r.predicted_answer for r in all_question_results]
        references = [r.expected_answer for r in all_question_results]
        questions = [r.question for r in all_question_results]

        # QA accuracy metrics (do NOT pass llm_judge here — scores already collected per-question)
        qa_metrics = compute_qa_metrics(
            predictions=predictions,
            references=references,
            fuzzy_threshold=self.fuzzy_threshold,
            llm_judge=None,
            questions=questions,
        )

        # Aggregate cached LLM judge scores if available
        llm_scores = [r.llm_judge_score for r in all_question_results if r.llm_judge_score is not None]
        if llm_scores:
            qa_metrics["avg_llm_judge_score"] = sum(llm_scores) / len(llm_scores)

        # MRR: treat fuzzy hit as binary relevance signal
        # Each query gets a ranked list of [predicted_answer]; if fuzzy hit → rank 1, else empty
        ranked_lists = []
        relevant_sets = []
        for r in all_question_results:
            if r.fuzzy_hit:
                ranked_lists.append([r.predicted_answer])
                relevant_sets.append({r.predicted_answer})
            else:
                ranked_lists.append([r.predicted_answer])
                relevant_sets.append({r.expected_answer})

        mrr = mean_reciprocal_rank(ranked_lists, relevant_sets)

        # Latency stats
        retrieve_latencies = [r.retrieve_latency_ms for r in all_question_results]
        remember_latencies = [sr.remember_latency_ms for sr in session_results]

        metrics = {
            **qa_metrics,
            "mrr": mrr,
            "total_questions": n,
            "total_sessions": len(session_results),
            "retrieve_latency": latency_stats(retrieve_latencies),
            "remember_latency": latency_stats(remember_latencies),
        }

        # Retrieval-style P/R/F1 if we treat exact hits as relevant
        # (over all questions combined)
        exact_predicted = [r.predicted_answer for r in all_question_results if r.exact_hit]
        exact_expected = [r.expected_answer for r in all_question_results]
        if exact_expected:
            prf = precision_recall_f1(
                [r.question_id for r in all_question_results if r.exact_hit],
                [r.question_id for r in all_question_results],
            )
            metrics["retrieval_precision"] = prf["precision"]
            metrics["retrieval_recall"] = prf["recall"]
            metrics["retrieval_f1"] = prf["f1"]

        return EvalReport(
            dataset_name=dataset_name,
            session_results=session_results,
            metrics=metrics,
            config={
                "fuzzy_threshold": self.fuzzy_threshold,
                "llm_judge": self.llm_judge is not None,
                "max_retrieved": self.max_retrieved,
            },
        )
