"""
cortex/eval/metrics.py — Scoring functions for memory eval benchmarks.

All functions work without an LLM (exact match, fuzzy match).
Optional LLM-as-judge is available when an llm_judge callable is provided.

Metrics:
  - exact_match(prediction, reference) -> bool
  - fuzzy_match(prediction, reference, threshold) -> bool
  - precision_recall_f1(predictions, references) -> dict
  - mean_reciprocal_rank(ranked_lists) -> float
  - latency_stats(latencies) -> dict
"""
from __future__ import annotations

import re
import time
from typing import Any, Callable, Dict, List, Optional, Sequence


# ---------------------------------------------------------------------------
# String normalisation
# ---------------------------------------------------------------------------

def _normalise(text: str) -> str:
    """Lowercase, strip punctuation/whitespace for comparison."""
    text = text.lower().strip()
    text = re.sub(r"[^\w\s]", "", text)
    text = re.sub(r"\s+", " ", text)
    return text.strip()


# ---------------------------------------------------------------------------
# Exact match
# ---------------------------------------------------------------------------

def exact_match(prediction: str, reference: str, normalise: bool = True) -> bool:
    """Return True if prediction exactly matches reference (optionally normalised)."""
    if normalise:
        return _normalise(prediction) == _normalise(reference)
    return prediction.strip() == reference.strip()


# ---------------------------------------------------------------------------
# Fuzzy match (token overlap / containment)
# ---------------------------------------------------------------------------

def fuzzy_match(
    prediction: str,
    reference: str,
    threshold: float = 0.6,
    method: str = "token_f1",
) -> bool:
    """
    Return True if fuzzy similarity between prediction and reference >= threshold.

    Methods:
      - token_f1: F1 on token overlap (default, no extra deps)
      - rapidfuzz: uses rapidfuzz.fuzz.partial_ratio if available
    """
    if method == "rapidfuzz":
        try:
            from rapidfuzz import fuzz
            score = fuzz.partial_ratio(_normalise(prediction), _normalise(reference)) / 100.0
            return score >= threshold
        except ImportError:
            method = "token_f1"  # fall back

    if method == "token_f1":
        pred_tokens = set(_normalise(prediction).split())
        ref_tokens = set(_normalise(reference).split())
        if not pred_tokens or not ref_tokens:
            return False
        overlap = pred_tokens & ref_tokens
        if not overlap:
            return False
        p = len(overlap) / len(pred_tokens)
        r = len(overlap) / len(ref_tokens)
        f1 = 2 * p * r / (p + r)
        return f1 >= threshold

    raise ValueError(f"Unknown fuzzy method: {method}")


def token_f1_score(prediction: str, reference: str) -> float:
    """Return token-level F1 score (0..1) between prediction and reference."""
    pred_tokens = _normalise(prediction).split()
    ref_tokens = _normalise(reference).split()
    if not pred_tokens or not ref_tokens:
        return 0.0
    pred_set = set(pred_tokens)
    ref_set = set(ref_tokens)
    overlap = pred_set & ref_set
    if not overlap:
        return 0.0
    p = len(overlap) / len(pred_set)
    r = len(overlap) / len(ref_set)
    return 2 * p * r / (p + r)


# ---------------------------------------------------------------------------
# LLM-as-judge (optional)
# ---------------------------------------------------------------------------

def llm_judge_score(
    prediction: str,
    reference: str,
    question: str,
    judge: Callable[[str, str, str], float],
) -> float:
    """
    Call an external LLM judge.

    Args:
        prediction: The system's answer.
        reference: The gold answer.
        question: The original question.
        judge: Callable(question, prediction, reference) -> float in [0, 1].

    Returns:
        Float score from 0 (wrong) to 1 (correct).
    """
    return judge(question, prediction, reference)


# ---------------------------------------------------------------------------
# Aggregate metrics
# ---------------------------------------------------------------------------

def compute_qa_metrics(
    predictions: List[str],
    references: List[str],
    fuzzy_threshold: float = 0.6,
    llm_judge: Optional[Callable[[str, str, str], float]] = None,
    questions: Optional[List[str]] = None,
) -> Dict[str, float]:
    """
    Compute QA-style metrics over a batch of predictions vs references.

    Returns a dict with:
      exact_match_accuracy, fuzzy_match_accuracy, avg_token_f1,
      (optionally) avg_llm_judge_score
    """
    if len(predictions) != len(references):
        raise ValueError(
            f"predictions ({len(predictions)}) and references ({len(references)}) must have same length"
        )

    n = len(predictions)
    if n == 0:
        return {
            "exact_match_accuracy": 0.0,
            "fuzzy_match_accuracy": 0.0,
            "avg_token_f1": 0.0,
            "n": 0,
        }

    exact_hits = 0
    fuzzy_hits = 0
    f1_scores: List[float] = []
    llm_scores: List[float] = []

    for i, (pred, ref) in enumerate(zip(predictions, references)):
        if exact_match(pred, ref):
            exact_hits += 1
        if fuzzy_match(pred, ref, threshold=fuzzy_threshold):
            fuzzy_hits += 1
        f1_scores.append(token_f1_score(pred, ref))

        if llm_judge is not None and questions is not None:
            q = questions[i] if i < len(questions) else ""
            llm_scores.append(llm_judge_score(pred, ref, q, llm_judge))

    result = {
        "exact_match_accuracy": exact_hits / n,
        "fuzzy_match_accuracy": fuzzy_hits / n,
        "avg_token_f1": sum(f1_scores) / n,
        "n": n,
    }

    if llm_scores:
        result["avg_llm_judge_score"] = sum(llm_scores) / len(llm_scores)

    return result


# ---------------------------------------------------------------------------
# Precision / Recall / F1 (retrieval-style)
# ---------------------------------------------------------------------------

def precision_recall_f1(
    retrieved_ids: Sequence[str],
    relevant_ids: Sequence[str],
) -> Dict[str, float]:
    """
    Compute precision, recall, and F1 for a single retrieval result.

    Args:
        retrieved_ids: IDs returned by the system.
        relevant_ids: Ground-truth relevant IDs.

    Returns:
        Dict with precision, recall, f1.
    """
    retrieved_set = set(retrieved_ids)
    relevant_set = set(relevant_ids)

    if not retrieved_set:
        return {"precision": 0.0, "recall": 0.0, "f1": 0.0}
    if not relevant_set:
        return {"precision": 0.0, "recall": 0.0, "f1": 0.0}

    tp = len(retrieved_set & relevant_set)
    precision = tp / len(retrieved_set)
    recall = tp / len(relevant_set)
    f1 = (2 * precision * recall / (precision + recall)) if (precision + recall) > 0 else 0.0

    return {"precision": precision, "recall": recall, "f1": f1}


def mean_reciprocal_rank(ranked_lists: List[List[str]], relevant_ids: List[Set[str]]) -> float:
    """
    Compute Mean Reciprocal Rank (MRR) over multiple queries.

    Args:
        ranked_lists: List of ranked result ID lists, one per query.
        relevant_ids: List of sets of relevant IDs, one per query.

    Returns:
        MRR score (0..1).
    """
    if len(ranked_lists) != len(relevant_ids):
        raise ValueError("ranked_lists and relevant_ids must have the same length")

    reciprocal_ranks = []
    for ranked, relevant in zip(ranked_lists, relevant_ids):
        rr = 0.0
        for rank, item in enumerate(ranked, start=1):
            if item in relevant:
                rr = 1.0 / rank
                break
        reciprocal_ranks.append(rr)

    if not reciprocal_ranks:
        return 0.0
    return sum(reciprocal_ranks) / len(reciprocal_ranks)


# Fix missing Set import
from typing import Set  # noqa: E402 (moved after definition for clarity)


# ---------------------------------------------------------------------------
# Latency stats
# ---------------------------------------------------------------------------

def latency_stats(latencies_ms: List[float]) -> Dict[str, float]:
    """
    Compute latency statistics from a list of measurements in milliseconds.

    Returns:
        Dict with mean, p50, p95, p99, min, max.
    """
    if not latencies_ms:
        return {"mean": 0.0, "p50": 0.0, "p95": 0.0, "p99": 0.0, "min": 0.0, "max": 0.0}

    sorted_lat = sorted(latencies_ms)
    n = len(sorted_lat)

    def percentile(p: float) -> float:
        idx = int(n * p / 100)
        return sorted_lat[min(idx, n - 1)]

    return {
        "mean": sum(sorted_lat) / n,
        "p50": percentile(50),
        "p95": percentile(95),
        "p99": percentile(99),
        "min": sorted_lat[0],
        "max": sorted_lat[-1],
        "count": n,
    }
