"""Second-pass verification ("counter run") over every generated digest document.

A synthesis pass can hallucinate a plausible-sounding quote, figure, or
theoretical connection even when its own prompt explicitly forbids it (see the
"No speculation" section of prompts/political.md, and similar language in the
other two prompts). Asking the same model "did you actually see this" in a
fresh call, with the source context in front of it again, catches a real
fraction of that which the generation prompt's instructions alone do not. This
is that second call — the final filter before a document is written to
output/, not a replacement for the instructions in the generation prompts.

Deliberately standalone rather than importing from llm_client.py: the two
modules would otherwise import each other (llm_client would need this module
to filter its own output; this module needs llm_client's Ollama-calling and
fallback-document helpers), and untangling that into a third shared module was
more churn than duplicating about a dozen lines of a plain HTTP POST.

Fails safe in the direction of showing LESS to the user, never more. Three
distinct failure modes, three distinct responses:
  - This pass cannot run at all (Ollama unreachable a second time, unparseable
    verdict) -> keep the original document, but mark it `unverified` rather
    than silently passing it through unchecked. Blanking the section entirely
    here would make a transient failure in THIS pass as damaging as the whole
    digest being down, which is a worse trade than a visibly-unverified
    document.
  - Specific windows/narration are flagged ungrounded -> drop just those, keep
    the rest.
  - Every window in a document is flagged -> withhold the section with an
    honest placeholder, the same shape generate_section() already uses when
    the LLM host doesn't respond at all.
"""

from __future__ import annotations

import json
import logging
import os
from datetime import datetime, timezone
from pathlib import Path

import requests

LOG = logging.getLogger(__name__)

PROMPT_PATH = Path(__file__).parent / "prompts" / "counter_run.md"


def _now_iso():
    return datetime.now(timezone.utc).replace(microsecond=0).isoformat().replace("+00:00", "Z")


def _withheld_document(document):
    section = document.get("section", "")
    return {
        "generated_at": _now_iso(),
        "detail_level": document.get("detail_level"),
        "section": section,
        "windows": [
            {
                "id": f"{section}-withheld",
                "title": f"{str(section).title()} — withheld pending verification",
                "kind": "text",
                "content": (
                    "Nothing in this run's draft for this section could be verified "
                    "against its sources, so it was withheld rather than shown "
                    "unverified."
                ),
            }
        ],
        "narration": "",
        "withheld": True,
    }


def _call_ollama(prompt, temperature):
    host = os.environ.get("OLLAMA_HOST", "http://llm-host:11434").rstrip("/")
    # A separate override, not a reuse of OLLAMA_MODEL by force: verification is a
    # different task from synthesis, and a household running a second, larger model
    # for it (or the same one) should be able to say so without the two being coupled.
    model = os.environ.get("COUNTER_RUN_MODEL") or os.environ.get("OLLAMA_MODEL", "qwen2.5:14b-instruct")
    timeout = float(os.environ.get("OLLAMA_TIMEOUT", "600"))

    payload = {
        "model": model,
        "prompt": prompt,
        "stream": False,
        "format": "json",
        "options": {"temperature": temperature},
    }
    response = requests.post(f"{host}/api/generate", json=payload, timeout=timeout)
    response.raise_for_status()
    return response.json().get("response", "")


def _build_prompt(document, context):
    prompt_text = PROMPT_PATH.read_text(encoding="utf-8")
    return (
        f"{prompt_text}\n\n"
        "## Source context\n\n```json\n"
        f"{json.dumps(context, indent=2, ensure_ascii=False, default=str)}\n"
        "```\n\n"
        "## Document to verify\n\n```json\n"
        f"{json.dumps(document, indent=2, ensure_ascii=False, default=str)}\n"
        "```\n"
    )


def _normalize(text):
    return " ".join(str(text).split()).strip().lower()


def _window_haystack(window):
    content = window.get("content")
    if isinstance(content, list):
        return _normalize(" ".join(str(item) for item in content))
    return _normalize(content or "")


def _ground_sources(container, context_blob):
    """Strips citations that don't trace back to the context, in place.

    A `sources` block is the strongest claim the digest makes — a URL and a quote
    together assert "this exists and says this" — and it is also the easiest thing
    for a model to compose out of thin air, because a plausible URL looks exactly
    like a real one. Both halves are checked mechanically here rather than being
    left to the verifying call's own judgement, for the same reason the quote check
    below exists: substring presence in the context is a fact, not an opinion, and
    no second opinion improves on it.

    A source with a fabricated URL is dropped whole. A real source carrying an
    invented quote keeps the source and loses the quote — the citation is still
    true, only the excerpt was not.
    """
    sources = container.get("sources")
    if not isinstance(sources, list):
        return

    kept = []
    for source in sources:
        if not isinstance(source, dict):
            continue
        url = str(source.get("url") or "").strip()
        if url and _normalize(url) not in context_blob:
            LOG.warning("counter-run: dropping a source whose URL is not in the context: %r", url[:120])
            continue
        quote = str(source.get("quote") or "").strip()
        if quote and _normalize(quote) not in context_blob:
            LOG.warning("counter-run: clearing a source quote not found in the context: %r", quote[:80])
            source = {key: value for key, value in source.items() if key != "quote"}
        kept.append(source)

    if kept:
        container["sources"] = kept
    else:
        container.pop("sources", None)


def verify_document(document, context):
    """Returns a possibly-filtered copy of `document`. Never raises."""
    if document.get("degraded"):
        # Already a system-generated "the LLM host did not respond" placeholder, not
        # model output — nothing here could be hallucinated, and the context it was
        # meant to be grounded in is exactly what it's reporting it never got.
        return document

    windows = document.get("windows") or []
    if not windows:
        return document

    temperature = float(os.environ.get("COUNTER_RUN_TEMPERATURE", "0.1"))

    try:
        raw = _call_ollama(_build_prompt(document, context), temperature)
        verdict = json.loads(raw)
        if not isinstance(verdict, dict):
            raise ValueError("verdict was not a JSON object")
    except Exception:
        LOG.warning(
            "counter-run: verification call failed for %s/%s, keeping the original "
            "document but marking it unverified",
            document.get("section"),
            document.get("detail_level"),
            exc_info=True,
        )
        result = dict(document)
        result["unverified"] = True
        return result

    window_verdicts = {
        str(entry.get("window_id")): bool(entry.get("grounded", True))
        for entry in (verdict.get("window_verdicts") or [])
        if isinstance(entry, dict)
    }

    # A quote the model claims is grounded gets checked mechanically too: this is
    # the one claim type a substring search can verify without another LLM call
    # trusting itself, so it's the one place a lenient (or dishonest) verdict from
    # the counter-run call itself still gets overridden.
    context_blob = _normalize(json.dumps(context, ensure_ascii=False, default=str))
    quote_failed_windows = set()
    for entry in verdict.get("quotes_checked") or []:
        if not isinstance(entry, dict):
            continue
        text = str(entry.get("text") or "")
        normalized = _normalize(text)
        if not normalized:
            continue
        actually_found = normalized in context_blob
        if entry.get("found_in_context", True) and not actually_found:
            LOG.warning(
                "counter-run: quote marked found_in_context but not located in "
                "context, overriding: %r",
                text[:80],
            )
        if not actually_found:
            for window in windows:
                if normalized in _window_haystack(window):
                    quote_failed_windows.add(str(window.get("id")))

    kept, dropped = [], []
    for window in windows:
        window_id = str(window.get("id"))
        grounded = window_verdicts.get(window_id, True) and window_id not in quote_failed_windows
        if grounded:
            kept.append(window)
        else:
            dropped.append(window_id)

    if dropped:
        LOG.warning(
            "counter-run: dropped %d ungrounded window(s) from %s/%s: %s",
            len(dropped),
            document.get("section"),
            document.get("detail_level"),
            dropped,
        )

    if not kept:
        return _withheld_document(document)

    # Citations are checked after the window verdicts, not before: there is no point
    # grounding the sources of a window that is about to be dropped whole. Markers
    # carry their own sources (the globe briefs), so they are walked too.
    for window in kept:
        _ground_sources(window, context_blob)
        for marker in window.get("globe_markers") or []:
            if isinstance(marker, dict):
                _ground_sources(marker, context_blob)

    result = dict(document)
    result["windows"] = kept
    if not bool(verdict.get("narration_grounded", True)):
        LOG.warning(
            "counter-run: narration for %s/%s flagged ungrounded, clearing it",
            document.get("section"),
            document.get("detail_level"),
        )
        result["narration"] = ""
    result["counter_run_checked"] = True
    return result


def verify_all(documents, section_contexts):
    """`documents` is {"compact": [doc, ...], "full": [doc, ...]} from
    llm_client.generate_all(); each doc's own `section` field looks up its context."""
    verified = {}
    for detail_level, docs in documents.items():
        verified[detail_level] = [
            verify_document(doc, section_contexts.get(doc.get("section"), {})) for doc in docs
        ]
    return verified