"""
cortex/eval/long_mt_bench.py — Long-MT-Bench+ benchmark adapter (#1281).

Long-MT-Bench+: Tests conversation continuation quality with long context.
Unlike LoCoMo/LongMemEval (which test stored memory retrieval), this benchmark
tests the ability to continue a conversation coherently given its full history.

Architecture:
    1. Load conversation history (many turns)
    2. Inject history into context (simulating memory injection)
    3. Present a follow-up question
    4. LLM-judge scores the response on coherence, factual consistency, and relevance

This is fundamentally different from stored-memory benchmarks:
    - No remember() call — history is directly injected
    - No retrieve() call — the full context is provided
    - Scoring is by LLM judge (1-10 scale, matching MT-Bench methodology)

Usage::

    from cortex.eval.long_mt_bench import LongMTBenchAdapter, LongMTBenchConfig

    adapter = LongMTBenchAdapter(llm=llm_provider)
    results = await adapter.evaluate(conversations)
"""
from __future__ import annotations

import json
import logging
import pathlib
import time
from dataclasses import dataclass, field
from typing import Any, Dict, List, Optional

logger = logging.getLogger(__name__)

_SAMPLE_PATH = pathlib.Path(__file__).parent / "datasets" / "long_mt_bench_sample.json"


# ---------------------------------------------------------------------------
# Data types
# ---------------------------------------------------------------------------

@dataclass
class MTBenchTurn:
    """A single conversation turn."""
    role: str  # 'user' or 'assistant'
    content: str


@dataclass
class MTBenchConversation:
    """A full conversation with a follow-up question for evaluation."""
    conversation_id: str
    history: List[MTBenchTurn]
    follow_up_question: str
    reference_answer: Optional[str] = None  # Gold answer for comparison
    category: str = "general"  # Topic category
    num_turns: int = 0

    def __post_init__(self):
        if self.num_turns == 0:
            self.num_turns = len(self.history)


@dataclass
class MTBenchResult:
    """Result for a single conversation evaluation."""
    conversation_id: str
    follow_up_question: str
    generated_answer: str
    judge_score: float  # 1-10 scale
    judge_reasoning: str
    category: str
    num_turns: int
    latency_ms: float
    error: Optional[str] = None


@dataclass
class MTBenchReport:
    """Aggregate results for the benchmark."""
    results: List[MTBenchResult] = field(default_factory=list)
    metrics: Dict[str, Any] = field(default_factory=dict)

    def summary(self) -> str:
        m = self.metrics
        lines = [
            "=== Long-MT-Bench+ Report ===",
            f"Conversations evaluated: {m.get('total', 0)}",
            f"Average judge score:     {m.get('avg_score', 0):.2f} / 10.0",
            f"Score >= 7 (good):       {m.get('good_rate', 0):.1%}",
            f"Score >= 5 (acceptable): {m.get('acceptable_rate', 0):.1%}",
        ]
        if "by_category" in m:
            lines.append("By category:")
            for cat, score in sorted(m["by_category"].items()):
                lines.append(f"  {cat}: {score:.2f}")
        return "\n".join(lines)

    def to_dict(self) -> Dict[str, Any]:
        return {
            "metrics": self.metrics,
            "results": [
                {
                    "conversation_id": r.conversation_id,
                    "follow_up_question": r.follow_up_question,
                    "generated_answer": r.generated_answer,
                    "judge_score": r.judge_score,
                    "judge_reasoning": r.judge_reasoning,
                    "category": r.category,
                    "num_turns": r.num_turns,
                    "latency_ms": r.latency_ms,
                    "error": r.error,
                }
                for r in self.results
            ],
        }


# ---------------------------------------------------------------------------
# Judge prompt (MT-Bench style)
# ---------------------------------------------------------------------------

_JUDGE_SYSTEM = """\
You are an expert judge evaluating AI assistant responses in long conversations.
Score the response on a scale of 1-10 based on:
- Coherence with conversation history
- Factual consistency with previously stated information
- Relevance and helpfulness to the follow-up question
- Quality of reasoning and detail

Output a JSON object with:
{
  "score": <1-10>,
  "reasoning": "<brief explanation>"
}
"""

_JUDGE_USER_TEMPLATE = """\
Here is a conversation history:

{history}

The user now asks:
{question}

The assistant responds:
{answer}

{reference_section}

Rate the response quality (1-10) based on coherence with history, factual consistency, and helpfulness.
"""


# ---------------------------------------------------------------------------
# Adapter
# ---------------------------------------------------------------------------

@dataclass
class LongMTBenchConfig:
    """Configuration for Long-MT-Bench+ evaluation."""
    max_history_tokens: int = 8000
    judge_model: Optional[str] = None
    generation_model: Optional[str] = None


class LongMTBenchAdapter:
    """Adapter for Long-MT-Bench+ evaluation.

    Unlike memory benchmarks, this adapter:
    1. Takes a conversation history
    2. Generates a response to a follow-up question with history as context
    3. Uses LLM-as-judge to score the response (1-10)
    """

    def __init__(
        self,
        llm: Any,
        config: Optional[LongMTBenchConfig] = None,
    ):
        self.llm = llm
        self.config = config or LongMTBenchConfig()

    async def evaluate(
        self,
        conversations: List[MTBenchConversation],
    ) -> MTBenchReport:
        """Evaluate a list of conversations."""
        results: List[MTBenchResult] = []

        for conv in conversations:
            result = await self._evaluate_conversation(conv)
            results.append(result)

        return self._build_report(results)

    async def _evaluate_conversation(
        self,
        conv: MTBenchConversation,
    ) -> MTBenchResult:
        """Evaluate a single conversation."""
        t0 = time.perf_counter()
        error = None

        # Format history
        history_text = self._format_history(conv.history)

        # Generate response
        try:
            gen_kwargs = {}
            if self.config.generation_model:
                gen_kwargs["model"] = self.config.generation_model

            prompt = f"""Continue this conversation naturally. Here is the history:

{history_text}

User: {conv.follow_up_question}
Assistant:"""

            resp = await self.llm.complete(
                system="You are a helpful assistant continuing a conversation. Use the history for context.",
                user=prompt,
                **gen_kwargs,
            )
            answer = resp.text if hasattr(resp, "text") else str(resp)
            answer = answer.strip()
        except Exception as exc:
            answer = ""
            error = str(exc)
            logger.warning("Generation failed for %s: %s", conv.conversation_id, exc)

        # Judge the response
        judge_score = 0.0
        judge_reasoning = ""
        if answer and not error:
            try:
                judge_score, judge_reasoning = await self._judge_response(
                    history_text, conv.follow_up_question, answer,
                    conv.reference_answer,
                )
            except Exception as exc:
                error = f"Judge failed: {exc}"
                logger.warning("Judge failed for %s: %s", conv.conversation_id, exc)

        latency_ms = (time.perf_counter() - t0) * 1000

        return MTBenchResult(
            conversation_id=conv.conversation_id,
            follow_up_question=conv.follow_up_question,
            generated_answer=answer,
            judge_score=judge_score,
            judge_reasoning=judge_reasoning,
            category=conv.category,
            num_turns=conv.num_turns,
            latency_ms=latency_ms,
            error=error,
        )

    async def _judge_response(
        self,
        history: str,
        question: str,
        answer: str,
        reference: Optional[str],
    ) -> tuple:
        """Judge a response using LLM-as-judge. Returns (score, reasoning)."""
        ref_section = ""
        if reference:
            ref_section = f"\nReference answer (for comparison): {reference}"

        prompt = _JUDGE_USER_TEMPLATE.format(
            history=history,
            question=question,
            answer=answer,
            reference_section=ref_section,
        )

        kwargs = {"response_format": "json", "max_tokens": 256, "temperature": 0.0}
        if self.config.judge_model:
            kwargs["model"] = self.config.judge_model

        resp = await self.llm.complete(
            system=_JUDGE_SYSTEM,
            user=prompt,
            **kwargs,
        )
        raw = resp.text if hasattr(resp, "text") else str(resp)

        try:
            data = json.loads(raw)
            score = float(data.get("score", 0))
            reasoning = data.get("reasoning", "")
            return min(10.0, max(1.0, score)), reasoning
        except (json.JSONDecodeError, ValueError, TypeError):
            # Fallback: try to extract a number
            import re
            m = re.search(r'\b(\d+(?:\.\d+)?)\b', raw)
            if m:
                return min(10.0, max(1.0, float(m.group(1)))), raw[:200]
            return 5.0, f"Parse error: {raw[:200]}"

    def _format_history(self, turns: List[MTBenchTurn]) -> str:
        """Format conversation history for prompt injection.

        Respects max_history_tokens by truncating from the beginning
        (keeping the most recent turns).
        """
        lines = []
        for turn in turns:
            role = turn.role.capitalize()
            lines.append(f"{role}: {turn.content}")

        text = "\n\n".join(lines)
        budget = self.config.max_history_tokens
        if budget and budget > 0:
            # Approximate tokens as words
            words = text.split()
            if len(words) > budget:
                text = " ".join(words[-budget:])
        return text

    def _build_report(self, results: List[MTBenchResult]) -> MTBenchReport:
        """Build aggregate report from individual results."""
        if not results:
            return MTBenchReport(metrics={"total": 0, "avg_score": 0.0})

        scores = [r.judge_score for r in results if r.error is None]
        n = len(scores)

        metrics: Dict[str, Any] = {
            "total": len(results),
            "evaluated": n,
            "avg_score": sum(scores) / n if n else 0.0,
            "good_rate": sum(1 for s in scores if s >= 7.0) / n if n else 0.0,
            "acceptable_rate": sum(1 for s in scores if s >= 5.0) / n if n else 0.0,
        }

        # By category
        cat_scores: Dict[str, List[float]] = {}
        for r in results:
            if r.error is None:
                cat_scores.setdefault(r.category, []).append(r.judge_score)
        if cat_scores:
            metrics["by_category"] = {
                cat: sum(ss) / len(ss) for cat, ss in sorted(cat_scores.items())
            }

        return MTBenchReport(results=results, metrics=metrics)


# ---------------------------------------------------------------------------
# Dataset loader
# ---------------------------------------------------------------------------

def load_long_mt_bench(
    path: Optional[str] = None,
) -> List[MTBenchConversation]:
    """Load Long-MT-Bench+ conversations from JSON.

    Expected format:
    [
      {
        "conversation_id": "...",
        "history": [{"role": "user", "content": "..."}, ...],
        "follow_up_question": "...",
        "reference_answer": "...",
        "category": "..."
      }
    ]
    """
    if path is None:
        path = str(_SAMPLE_PATH)

    p = pathlib.Path(path)
    if not p.exists():
        raise FileNotFoundError(f"Long-MT-Bench+ dataset not found: {p}")

    with open(p, "r", encoding="utf-8") as f:
        raw = json.load(f)

    if isinstance(raw, dict):
        if "conversations" in raw:
            raw = raw["conversations"]
        elif "data" in raw:
            raw = raw["data"]
        else:
            raise ValueError(
                f"Long-MT-Bench+ dict must have 'conversations' or 'data' key, "
                f"got keys: {list(raw.keys())}"
            )

    conversations: List[MTBenchConversation] = []
    for item in raw:
        history = [
            MTBenchTurn(role=t.get("role", "user"), content=t.get("content", ""))
            for t in item.get("history", item.get("turns", []))
        ]
        conversations.append(
            MTBenchConversation(
                conversation_id=item.get("conversation_id", item.get("id", "")),
                history=history,
                follow_up_question=item.get("follow_up_question", item.get("question", "")),
                reference_answer=item.get("reference_answer", item.get("answer")),
                category=item.get("category", "general"),
            )
        )

    return conversations
