#!/usr/bin/env python3
"""Multi-domain dogfood: prove A (faithfulness) + B (completeness) beyond the Ambipar corpus.

Per domain, against a PUBLIC rules document:
  1. ingest  — chunk + embed the document into the local engine store (source-scoped).
  2. report  — have the routed LLM write a ~20-claim analytical memo FROM THE FULL DOCUMENT
               (the customer's side of the loop, generated honestly, no retrieval involved).
  3. corrupt — have the LLM introduce exactly K factual value errors into that memo and
               DECLARE them (original vs corrupted), so ground truth is known.
  4. review  — run PortMem.review() (A+B) on the clean memo (expect high score, no false
               contradictions) and on the corrupted memo (expect the K errors flagged).
  5. score   — per seeded error, find the claim carrying the corrupted value: verdict
               contradicted/uncertain = CAUGHT; supported = MISSED (the unsafe error).

Artifacts land in eval/dogfood/<source-id>/. Uses the local docker Postgres
(portmem-engine-pg, :5434) + config/routing.haiku.yaml (prod-parity judge: Haiku 4.5;
Qwen3-8B embeddings matching the store's 4096d vectors).

  python3 dogfood_domain.py --doc eval/dogfood/corpora/x.txt --source-id "FDA Label X" --stage all
"""
from __future__ import annotations

import argparse
import asyncio
import json
import os
import re
import time
from pathlib import Path

from run_local_verify import _load_env

ROOT = Path(__file__).resolve().parent.parent
OUT = ROOT / "eval" / "dogfood"

ENV_DEFAULTS = {
    "MDMA_POSTGRES_DSN": "postgresql+asyncpg://mdma:mdma_dev_password@localhost:5434/mdma",
    "MDMA_USE_PG_GRAPH": "1",
    "MDMA_USE_PG_MEMORY": "1",
    "MDMA_MODEL_ROUTING_CONFIG_PATH": str(ROOT / "config" / "routing.haiku.yaml"),
}

REPORT_PROMPT = """You are a diligent analyst. Write an analytical memo summarizing the key
provisions of the document below for a professional audience. State roughly 20 specific,
checkable facts (figures, dates, thresholds, named parties, obligations). Use only what the
document says; do not add outside knowledge. Plain prose paragraphs and bullet lists, no
preamble, no closing boilerplate.

DOCUMENT:
{doc}"""

CORRUPT_PROMPT = """Take the memo below and introduce EXACTLY {k} factual value errors: change a
number, date, percentage, threshold, duration, or named entity to a plausible but WRONG value.
Keep everything else verbatim. Spread the errors across different parts of the memo.

Return ONLY JSON:
{{"report": "<the full memo with the {k} errors>",
  "errors": [{{"original": "<exact original phrase>", "corrupted": "<exact corrupted phrase>"}}, ...]}}

MEMO:
{memo}"""


def _dirs(source_id: str) -> Path:
    d = OUT / re.sub(r"[^A-Za-z0-9._-]+", "_", source_id)
    d.mkdir(parents=True, exist_ok=True)
    return d


async def _pm():
    from portmem.config import PortmemConfig
    from portmem.engine_binding import PortMem
    cfg = PortmemConfig(deployment={"mode": "cloud"}, capabilities={})
    return await PortMem.from_config(cfg)


async def _llm(pm, prompt: str, max_tokens: int = 4096) -> str:
    from mdma.router.task_types import TaskType
    router = pm.backend._state.router
    entry = router._routing[TaskType.SUMMARIZATION].chain[0]
    provider = router._registry.get_llm(entry.provider_name)
    r = await provider.complete(prompt=prompt, model=entry.model, max_tokens=max_tokens,
                                temperature=0.2, timeout_seconds=300)
    return r.text


def _first_json(text: str) -> dict:
    m = re.search(r"\{.*\}", text, re.DOTALL)
    return json.loads(m.group(0))


def _diff_tokens(original: str, corrupted: str) -> list[str]:
    """Tokens present in the corrupted phrase but NOT the original — the actual
    injected difference. This is what a claim must carry to count as 'the seeded
    claim'; matching on any number in the phrase (the old rule) degenerated to
    fragments like '1' and credited one contradicted claim to three unrelated
    seeds (2026-08-08 eval audit)."""
    tok = lambda s: re.findall(r"[A-Za-z]+|[\d][\d.,%x/-]*", s.lower())
    orig = set(tok(original))
    return [t for t in tok(corrupted) if t not in orig]


def _claim_matches(claim: str, e: dict) -> bool:
    """A claim carries the seeded error iff it contains a genuine diff token on a
    word boundary AND shares enough of the corrupted sentence's vocabulary that it
    is plausibly the same statement (guards against a diff word like 'three'
    matching an unrelated sentence)."""
    c = claim.lower()
    diffs = _diff_tokens(e.get("original", ""), e["corrupted"])
    if not diffs:
        diffs = [e["corrupted"].lower()[:40]]
    if not any(re.search(rf"(?<![\w]){re.escape(d)}(?![\w])", c) for d in diffs):
        return False
    ctx = [t for t in re.findall(r"[a-z]{4,}", e["corrupted"].lower())]
    if not ctx:
        return True
    overlap = sum(1 for t in set(ctx) if t in c)
    return overlap >= max(1, min(3, len(set(ctx)) // 3))


def score_seeded(review: dict, errors: list[dict]) -> list[dict]:
    """Map each seeded error to the reviewed claim carrying its corrupted value.

    Matching rule (2026-08-08, post-audit): a claim matches a seed only if it
    carries the original->corrupted DIFF token on a word boundary and shares
    vocabulary with the corrupted sentence. `caught` = any matching claim came
    back contradicted/uncertain; `false_supported` = matching claims exist and
    ALL came back supported (the unsafe outcome); `unmatched` = the corrupted
    statement was never extracted as a claim, i.e. never reviewed — not a miss,
    but excluded from any 'of errors that entered review' numerator AND
    denominator. A mixed supported+uncertain set is NOT false_supported because
    the flagged verdict routes the claim to human review; per-seed verdict lists
    are recorded so that call is auditable."""
    claims = review["faithfulness"]["claims"]
    out = []
    for e in errors:
        hits = [c for c in claims if _claim_matches(c["claim"], e)]
        verdicts = sorted({c["verdict"] for c in hits})
        caught = any(v in ("contradicted", "uncertain") for v in verdicts)
        missed_supported = bool(hits) and all(v == "supported" for v in verdicts)
        out.append({**e, "matched_claims": len(hits), "verdicts": verdicts,
                    "matched_texts": [c["claim"][:120] for c in hits],
                    "caught": caught, "false_supported": missed_supported,
                    "unmatched": not hits})
    return out


async def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("--doc", required=True)
    ap.add_argument("--source-id", required=True)
    ap.add_argument("--stage", default="all", choices=["all", "ingest", "report", "review"])
    ap.add_argument("--k-errors", type=int, default=6)
    ap.add_argument("--max-claims", type=int, default=25)
    a = ap.parse_args()

    # Harness defaults BEFORE .env: both use setdefault, so precedence becomes
    # shell exports > these local-stack defaults > .env (whose routing/DSN may
    # point at another deployment; this harness is defined against the local
    # docker store + prod-parity Haiku routing).
    for k, v in ENV_DEFAULTS.items():
        os.environ.setdefault(k, v)
    _load_env()

    d = _dirs(a.source_id)
    doc = open(a.doc).read()
    pm = await _pm()
    print(f"[{a.source_id}] doc={len(doc)} chars -> {d}", flush=True)

    if a.stage in ("all", "ingest"):
        t0 = time.monotonic()
        r = await pm.ingest_source(doc, source_id=a.source_id)
        print(f"ingested: {r.get('chunks', r)} chunks in {time.monotonic()-t0:.0f}s", flush=True)

    clean_p, corr_p, err_p = d / "report_clean.txt", d / "report_corrupted.txt", d / "errors.json"
    if a.stage in ("all", "report"):
        memo = await _llm(pm, REPORT_PROMPT.format(doc=doc[:400_000]))
        clean_p.write_text(memo)
        c = _first_json(await _llm(pm, CORRUPT_PROMPT.format(k=a.k_errors, memo=memo), 8192))
        corr_p.write_text(c["report"])
        err_p.write_text(json.dumps(c["errors"], indent=1, ensure_ascii=False))
        print(f"memo={len(memo)} chars, seeded {len(c['errors'])} errors", flush=True)

    if a.stage in ("all", "review"):
        for label, path in (("clean", clean_p), ("corrupted", corr_p)):
            t0 = time.monotonic()
            rv = await pm.review(path.read_text(), max_claims=a.max_claims, max_results=8,
                                 source=a.source_id, check_completeness=(label == "clean"))
            f = rv["faithfulness"]
            line = (f"[{label}] score={f.get('score')} "
                    f"({f.get('n_supported')}sup/{f.get('n_contradicted')}contra/"
                    f"{f.get('n_uncertain')}unc/{f.get('n_not_found')}nf) "
                    f"{time.monotonic()-t0:.0f}s")
            if label == "corrupted":
                seeded = score_seeded(rv, json.loads(err_p.read_text()))
                rv["seeded_scoring"] = seeded
                caught = sum(s["caught"] for s in seeded)
                fs = sum(s["false_supported"] for s in seeded)
                um = sum(s["unmatched"] for s in seeded)
                line += f" | seeded: {caught}/{len(seeded)} caught, {fs} false-supported, {um} unmatched"
            cmp_ = rv.get("completeness")
            if cmp_:
                line += f" | B: {cmp_['n_covered']}/{cmp_['n_provisions']} covered, {cmp_['n_omissions']} omissions"
            print(line, flush=True)
            json.dump(rv, open(d / f"review_{label}.json", "w"), indent=1, default=str)
    print("done", flush=True)


if __name__ == "__main__":
    asyncio.run(main())