"""
LLMJudge — LLM-as-judge pseudo-labeling for borderline chains.

Problem
-------
LabelFreeScorer needs labeled chains for supervised calibration (QARA-Obs),
but human labeling is expensive and slow. For borderline chains (risk 0.45–0.65)
the model is uncertain — these are exactly the chains that would most benefit
from a correct label.

Solution
--------
Use a cheap LLM (Claude Haiku or GPT-3.5-mini) to classify whether the agent
chain produced a correct final answer. The LLM judges based on:
  1. The original question
  2. The agent's final answer
  3. (Optional) A reference answer or retrieved context

The pseudo-label confidence is calibrated against human labels:
  - Haiku accuracy on TriviaQA: ~89% (exp41_llm_judge results)
  - GPT-3.5 accuracy: ~85% estimated
  - Use confidence_threshold to only keep high-confidence labels

Integration with LabelFreeScorer
---------------------------------
    from qppg_service.llm_judge import LLMJudge, JudgeLabel
    from qppg_service.label_free_scorer import LabelFreeScorer

    judge  = LLMJudge(model="claude-haiku-4-5-20251001")
    scorer = LabelFreeScorer()

    # Pseudo-label a batch of borderline chains
    chains_to_label = [c for c in recent_chains if 0.45 <= c["risk_score"] <= 0.65]
    for chain in chains_to_label:
        label = judge.label(chain["question"], chain["steps"], chain["final_answer"])
        if label.confidence >= 0.80:
            chain["correct"] = label.correct

    # Use pseudo-labels for supervised calibration
    scorer.calibrate_supervised(chains_to_label)

Research context
----------------
  Exp 41 (llm_judge.py) explored this with HotpotQA chains:
    - Haiku judge accuracy: ~89% on within-domain chains
    - Cross-domain accuracy: ~82%
    - Best use case: chains with clear final answers (not hedged)
    - Worst case: multi-hop questions where Haiku lacks world knowledge

  Expected AUROC improvement from pseudo-labels:
    - Within-domain: AUROC 0.883 → 0.895+ (marginal, already high)
    - Cross-domain: AUROC 0.664 → 0.720+ (significant gain)
"""

from __future__ import annotations

import hashlib
import json
import time
from dataclasses import dataclass
from pathlib import Path
from typing import Optional


@dataclass
class JudgeLabel:
    """Result of LLM judgment on a single chain."""
    correct:     bool           # True = agent answered correctly
    confidence:  float          # [0, 1]; 1 = very confident
    reasoning:   str            # one-sentence explanation
    model:       str            # model that produced the label
    cached:      bool = False   # True if loaded from disk cache
    latency_ms:  float = 0.0


_JUDGE_PROMPT = """\
You are evaluating whether an AI agent answered a question correctly.

Question: {question}

Agent's final answer: {final_answer}

{ref_block}

Evaluate whether the agent's answer is factually correct and directly answers
the question. Respond with EXACTLY this JSON format (no markdown, no extra text):
{{"correct": true/false, "confidence": 0.0-1.0, "reasoning": "one sentence"}}

Rules:
- correct=true if the answer is factually right and answers the question
- correct=false if wrong, hallucinated, off-topic, or "I don't know"
- confidence reflects how sure YOU are (not the agent)
- reasoning must be under 20 words"""


class LLMJudge:
    """
    LLM-as-judge for pseudo-labeling agent chains.

    Parameters
    ----------
    model            : Claude or OpenAI model identifier
    confidence_threshold : minimum confidence to accept a pseudo-label
    cache_dir        : disk path for caching responses (saves API cost)
    max_retries      : number of retries on API errors
    """

    def __init__(
        self,
        model: str = "claude-haiku-4-5-20251001",
        confidence_threshold: float = 0.75,
        cache_dir: Optional[str | Path] = None,
        max_retries: int = 2,
    ):
        self.model                = model
        self.confidence_threshold = confidence_threshold
        self.max_retries          = max_retries
        self._cache: dict[str, dict] = {}

        self._cache_dir = Path(cache_dir).expanduser() if cache_dir else None
        if self._cache_dir:
            self._cache_dir.mkdir(parents=True, exist_ok=True)

    # ── Public API ────────────────────────────────────────────────────────────

    def label(
        self,
        question:     str,
        steps:        list[dict],
        final_answer: str,
        reference:    str = "",
    ) -> JudgeLabel:
        """
        Judge whether the agent's final answer is correct.

        Parameters
        ----------
        question     : original user question
        steps        : agent steps (used to extract context, not directly judged)
        final_answer : the agent's concluding answer
        reference    : (optional) gold/reference answer for calibration
        """
        cache_key = self._cache_key(question, final_answer, reference)
        cached    = self._load_cache(cache_key)
        if cached:
            return JudgeLabel(**{**cached, "cached": True})

        prompt  = self._build_prompt(question, final_answer, reference)
        t0      = time.time()
        raw     = self._call_llm(prompt)
        latency = (time.time() - t0) * 1000

        label = self._parse_response(raw, latency)
        self._save_cache(cache_key, label)
        return label

    def label_batch(
        self,
        chains:          list[dict],
        borderline_only: bool = True,
        scorer=None,
    ) -> list[Optional[JudgeLabel]]:
        """
        Label a batch of chain dicts.

        Parameters
        ----------
        chains          : list of chain dicts (must have question, steps, final_answer)
        borderline_only : if True and scorer is given, only label chains with
                          risk in [0.45, 0.65] (saves API calls for clear cases)
        scorer          : LabelFreeScorer for filtering borderline chains

        Returns a list of JudgeLabel (or None if chain was skipped).
        """
        labels = []
        for chain in chains:
            skip = False
            if borderline_only and scorer is not None:
                r = scorer.score(
                    chain["question"], chain["steps"],
                    chain.get("final_answer", ""), chain.get("finished", True)
                )
                if r.risk_score < 0.40 or r.risk_score > 0.70:
                    skip = True

            if skip:
                labels.append(None)
            else:
                lbl = self.label(
                    chain["question"],
                    chain["steps"],
                    chain.get("final_answer", ""),
                    chain.get("gold_answer", ""),
                )
                labels.append(lbl)

        return labels

    def filter_confident(
        self,
        chains: list[dict],
        labels: list[Optional[JudgeLabel]],
    ) -> list[dict]:
        """
        Return chains that have confident pseudo-labels.
        Sets chain["correct"] = label.correct for downstream calibration.
        """
        confident = []
        for chain, label in zip(chains, labels):
            if label is not None and label.confidence >= self.confidence_threshold:
                chain = dict(chain)
                chain["correct"] = label.correct
                confident.append(chain)
        return confident

    # ── LLM call ─────────────────────────────────────────────────────────────

    def _build_prompt(self, question: str, final_answer: str, reference: str) -> str:
        ref_block = ""
        if reference:
            ref_block = f"Reference answer: {reference}\n\n"
        return _JUDGE_PROMPT.format(
            question=question[:500],
            final_answer=(final_answer or "[no answer]")[:500],
            ref_block=ref_block,
        )

    def _call_llm(self, prompt: str) -> str:
        """Call the configured LLM. Supports Claude and OpenAI."""
        is_claude = "claude" in self.model.lower()

        for attempt in range(self.max_retries + 1):
            try:
                if is_claude:
                    return self._call_claude(prompt)
                else:
                    return self._call_openai(prompt)
            except Exception as e:
                if attempt == self.max_retries:
                    raise
                time.sleep(2 ** attempt)
        raise RuntimeError("Unreachable")

    def _call_claude(self, prompt: str) -> str:
        import anthropic  # type: ignore
        client = anthropic.Anthropic()
        msg    = client.messages.create(
            model      = self.model,
            max_tokens = 150,
            messages   = [{"role": "user", "content": prompt}],
        )
        return msg.content[0].text.strip()

    def _call_openai(self, prompt: str) -> str:
        import openai  # type: ignore
        client = openai.OpenAI()
        resp   = client.chat.completions.create(
            model       = self.model,
            max_tokens  = 150,
            messages    = [{"role": "user", "content": prompt}],
            temperature = 0.0,
        )
        return resp.choices[0].message.content.strip()

    def _parse_response(self, raw: str, latency_ms: float) -> JudgeLabel:
        try:
            # Strip markdown code fences if present
            text = raw.strip()
            if text.startswith("```"):
                text = text.split("```")[1]
                if text.lower().startswith("json"):
                    text = text[4:]
            d = json.loads(text)
            return JudgeLabel(
                correct    = bool(d.get("correct", False)),
                confidence = float(d.get("confidence", 0.5)),
                reasoning  = str(d.get("reasoning", ""))[:200],
                model      = self.model,
                latency_ms = latency_ms,
            )
        except Exception:
            # Fallback: look for true/false in raw text
            correct    = "true" in raw.lower() and "false" not in raw.lower()[:raw.lower().index("true")]
            return JudgeLabel(
                correct    = correct,
                confidence = 0.5,
                reasoning  = "parse error — low confidence fallback",
                model      = self.model,
                latency_ms = latency_ms,
            )

    # ── Cache ─────────────────────────────────────────────────────────────────

    def _cache_key(self, question: str, final_answer: str, reference: str) -> str:
        s = json.dumps([question[:300], final_answer[:300], reference[:100]])
        return hashlib.md5(s.encode()).hexdigest()

    def _load_cache(self, key: str) -> Optional[dict]:
        if key in self._cache:
            return self._cache[key]
        if self._cache_dir:
            p = self._cache_dir / f"{key}.json"
            if p.exists():
                d = json.loads(p.read_text())
                self._cache[key] = d
                return d
        return None

    def _save_cache(self, key: str, label: JudgeLabel) -> None:
        d = {
            "correct":    label.correct,
            "confidence": label.confidence,
            "reasoning":  label.reasoning,
            "model":      label.model,
            "latency_ms": label.latency_ms,
        }
        self._cache[key] = d
        if self._cache_dir:
            p = self._cache_dir / f"{key}.json"
            p.write_text(json.dumps(d))
