#!/usr/bin/env python3
"""invoice-extract — dual-run orchestrator (model_policy in code).

Per document: deterministic pre-screens → two independent extraction legs
(M4 local MLX ∥ Mistral via vibe) → validators on each leg → exact agreement
required on critical fields. Disagreement, single-valid-leg, or an escalate
flag, or both legs invalid → the item lands in escalations/ for the Claude
tier; hostile documents quarantine at the pre-screen. Refuse, never repair.

Usage:
  dual_run.py --inputs <dir> [--injection <dir>] --out <predictions-dir> \\
      [--mlx-model ID] [--journal FILE] [--only id1,id2]

Predictions dir gets one <id>.json per resolved item (extraction payload or
{"outcome": "quarantine", ...}); unresolved disagreements go to
<out>/../escalations/<id>.json and are NOT written to the predictions dir —
score.py then reports them missing, which is the honest state until the
escalation tier resolves them.
"""
from __future__ import annotations

import argparse
import glob
import json
import os
import sys
import time

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
import extract
import validators

CRITICAL = ("totals.ht", "totals.tva", "totals.ttc", "ref_supplier", "date_issue", "date_due", "iban")


def _get(payload: dict, dotted: str):
    cur = payload
    for part in dotted.split("."):
        if not isinstance(cur, dict):
            return None
        cur = cur.get(part)
    return cur


def criticals_agree(a: dict, b: dict) -> list[str]:
    """Return the critical fields on which the two legs disagree."""
    diff = []
    for f in CRITICAL:
        va, vb = _get(a, f), _get(b, f)
        if f.startswith("totals."):
            same = va is not None and vb is not None and abs(float(va) - float(vb)) <= 0.001
        elif f == "iban":
            norm = lambda v: None if v in (None, "") else str(v).replace(" ", "").upper()
            same = norm(va) == norm(vb)
        else:
            same = va == vb
        if not same:
            diff.append(f"{f}: {va!r} vs {vb!r}")
    return diff


def process(item_id: str, doc: dict, mlx_model: str, journal) -> tuple[str, dict]:
    """Returns (disposition, record): disposition ∈ predict|quarantine|escalate."""
    text, sha = doc["text"], doc["source_sha256"]
    reasons, flags = validators.screen_document(text)
    if reasons:
        return "quarantine", {"outcome": "quarantine", "stage": "pre-screen", "why": reasons}

    legs = []
    for runtime, model in (("mlx", mlx_model), ("vibe", None)):
        t0 = time.time()
        try:
            leg = extract.run_leg(text, runtime, model)
        except Exception as e:  # endpoint down, timeout — an honest leg failure
            leg = {"runtime": runtime, "model": model, "ok": False, "error": str(e)[:300]}
        leg["latency_s"] = round(time.time() - t0, 1)
        if leg.get("ok"):
            validated, vreasons = validators.check(leg["payload"], text, sha)
            leg["valid"] = validated is not None
            leg["validated"] = validated
            leg["invariant_failures"] = vreasons
        else:
            leg["valid"] = False
        legs.append(leg)
        journal.write(json.dumps({"id": item_id, "leg": {k: v for k, v in leg.items() if k != "validated"}},
                                 ensure_ascii=False) + "\n")
        journal.flush()

    a, b = legs
    if not a["valid"] and not b["valid"]:
        # Hostile content is caught pre-screen; a clean document both legs fail
        # to ground goes to the escalation tier, whose verdict may be quarantine.
        return "escalate", {"stage": "escalation", "reason": "both-legs-invalid",
                            "legs": [{k: l.get(k) for k in ("runtime", "model", "valid",
                                                            "invariant_failures", "error")} for l in legs]}
    if flags or not (a["valid"] and b["valid"]):
        return "escalate", {"stage": "escalation", "flags": flags,
                            "legs": [{k: l.get(k) for k in ("runtime", "model", "valid", "validated",
                                                            "invariant_failures", "error")} for l in legs]}
    diff = criticals_agree(a["validated"], b["validated"])
    if diff:
        return "escalate", {"stage": "escalation", "flags": flags, "disagreement": diff,
                            "legs": [{k: l.get(k) for k in ("runtime", "model", "validated")} for l in legs]}
    out = dict(a["validated"])
    out["confidence"] = 0.9  # dual-leg exact agreement on criticals
    return "predict", out


def main() -> int:
    ap = argparse.ArgumentParser()
    ap.add_argument("--inputs", required=True)
    ap.add_argument("--injection")
    ap.add_argument("--out", required=True)
    ap.add_argument("--mlx-model", default="leonsarmiento/Ornith-1.0-35B-5bit-mlx")
    ap.add_argument("--journal", default=None)
    ap.add_argument("--only", default=None)
    args = ap.parse_args()

    os.makedirs(args.out, exist_ok=True)
    esc_dir = os.path.join(os.path.dirname(os.path.abspath(args.out)), "escalations")
    os.makedirs(esc_dir, exist_ok=True)
    only = set(args.only.split(",")) if args.only else None

    files = sorted(glob.glob(os.path.join(args.inputs, "*.json")))
    if args.injection:
        files += sorted(glob.glob(os.path.join(args.injection, "*.json")))

    journal_path = args.journal or os.path.join(os.path.dirname(args.out), "dual-run-journal.jsonl")
    counts = {"predict": 0, "quarantine": 0, "escalate": 0}
    with open(journal_path, "a") as journal:
        for f in files:
            item_id = os.path.splitext(os.path.basename(f))[0]
            if only and item_id not in only:
                continue
            doc = json.load(open(f))
            disposition, record = process(item_id, doc, args.mlx_model, journal)
            counts[disposition] += 1
            dest = os.path.join(args.out if disposition != "escalate" else esc_dir, item_id + ".json")
            json.dump(record, open(dest, "w"), indent=2, ensure_ascii=False)
            print(f"{item_id}: {disposition}", flush=True)
    print(json.dumps(counts))
    return 0


if __name__ == "__main__":
    sys.exit(main())