"""LLM entailment judge for ``retrievalExtract_contains`` assertions.

Decides whether an assertion's text is *supported by* a piece of retrieved
evidence, returning a ``{"supported": bool, "reason": str}`` verdict. The
judgment is delegated to an Azure OpenAI model via a prompty flow loaded with
the ``load_flow`` pattern established by the PII evaluator
(``custom_evaluators/pii/pii_evaluator.py``); ambiguous or unparseable model output
fails closed (``supported=False``).

This replaces V1's deterministic case-insensitive substring containment, which
produced false positives ("$500" matched "$500K") and false negatives
("out-of-pocket" vs "out of pocket"). For ``retrievalExtract_contains``, this
supersedes ``adr-string-matching-vs-llm-judge.md``.
"""

from __future__ import annotations

import json
import os
import re
from typing import Any, Dict, Optional

from promptflow.client import load_flow

_PROMPTY_FILENAME = "extract_judge.prompty"

# Strips a markdown code fence (```json ... ``` or ``` ... ```) so JSON wrapped
# in a fence - a common model habit even at temperature 0 - still parses.
_CODE_FENCE_RE = re.compile(r"```(?:json)?\s*(.*?)\s*```", re.IGNORECASE | re.DOTALL)

# Verdict/reason tag fallback when the model does not return clean JSON.
_VERDICT_RE = re.compile(r"<verdict>\s*(supported|not_supported)\s*</verdict>", re.IGNORECASE)
_REASON_RE = re.compile(r"<reason>(.*?)</reason>", re.IGNORECASE | re.DOTALL)

# String renderings of a truthy ``supported`` verdict; anything else is False.
_TRUE_STRINGS = frozenset({"true", "yes", "supported", "1"})

# Cap every reason surfaced to the report so a rambling model
# response can't bloat the report's reason field.
_MAX_REASON_CHARS = 500


class ExtractJudge:
    """Judge whether a single extract assertion is supported by retrieved evidence."""

    def __init__(self, model_config: Any) -> None:
        if model_config is None:
            raise ValueError(
                "ExtractJudge requires an Azure OpenAI model configuration; "
                "retrievalExtract_contains assertions need Azure OpenAI configured."
            )
        prompty_path = os.path.join(os.path.dirname(__file__), _PROMPTY_FILENAME)
        self._flow = load_flow(source=prompty_path, model={"configuration": model_config})

    def judge(self, *, claim: str, evidence: str) -> Dict[str, Any]:
        """Return ``{"supported": bool, "reason": str}`` for one claim vs. evidence.

        Never raises: a model call failure fails closed with the error in
        ``reason`` so one flaky call can't abort the whole evaluation.
        """
        try:
            raw = self._flow(claim=claim, evidence=evidence)
        except Exception as exc:
            return {
                "supported": False,
                "reason": _clip_reason(f"Extract judge error: {exc}"),
            }
        return _parse_verdict(raw)


def _parse_verdict(raw: Any) -> Dict[str, Any]:
    """Parse the model output. Prefers JSON, falls back to verdict/reason tags.

    Always returns a dict with bool ``supported`` and str ``reason``. An
    unparseable response is treated as ``supported=False`` so an ambiguous judge
    response fails closed rather than silently passing an assertion.

    JSON is extracted leniently: bare JSON, JSON in a markdown code fence, and
    a JSON object embedded in prose are all accepted, since models routinely
    add fences or prose despite the prompt.
    """
    if isinstance(raw, dict):
        return _verdict_from_mapping(raw)

    if not isinstance(raw, str):
        return {"supported": False, "reason": ""}

    data = _extract_json_object(raw)
    if data is not None:
        return _verdict_from_mapping(data)

    verdict = _VERDICT_RE.search(raw)
    reason = _REASON_RE.search(raw)
    reason_text = reason.group(1) if reason else raw
    return {
        "supported": bool(verdict) and verdict.group(1).lower() == "supported",
        "reason": _clip_reason(reason_text.strip()),
    }


def _clip_reason(reason: str) -> str:
    """Bound a reason to ``_MAX_REASON_CHARS`` so a verbose model response - on
    any path (JSON, tag, or raw fallback) - can't bloat the report."""
    return reason[:_MAX_REASON_CHARS]


def _verdict_from_mapping(data: Dict[str, Any]) -> Dict[str, Any]:
    """Normalize a parsed verdict mapping into the canonical verdict dict."""
    return {
        "supported": _coerce_supported(data.get("supported")),
        "reason": _clip_reason(str(data.get("reason", "")).strip()),
    }


def _coerce_supported(value: Any) -> bool:
    """Interpret a ``supported`` value, tolerating stringified booleans.

    ``bool("false")`` is ``True`` in Python, so a model that emits
    ``{"supported": "false"}`` must not be read as supported. Strings are
    matched case-insensitively against a small truthy set; everything else uses
    plain truthiness.
    """
    if isinstance(value, bool):
        return value
    if isinstance(value, str):
        return value.strip().casefold() in _TRUE_STRINGS
    return bool(value)


def _extract_json_object(raw: str) -> Optional[Dict[str, Any]]:
    """Best-effort parse of a JSON object from possibly-noisy model output.

    Tries, in order: the whole (stripped) string, the contents of a markdown
    code fence, and the substring spanning the first ``{`` to the last ``}``
    (an object embedded in prose). Returns the first candidate that parses to a
    dict, or ``None`` when none do.
    """
    candidates = [raw.strip()]

    fence = _CODE_FENCE_RE.search(raw)
    if fence:
        candidates.append(fence.group(1).strip())

    start = raw.find("{")
    end = raw.rfind("}")
    if start != -1 and end > start:
        candidates.append(raw[start:end + 1])

    for candidate in candidates:
        try:
            data = json.loads(candidate)
        except Exception:
            continue
        if isinstance(data, dict):
            return data
    return None