"""Run the pilot detectors over the pilot items. Resumable; one JSON line per item.

    python experiments/P0-pilot/run_pilot.py --items experiments/P0-pilot/data/items.jsonl \
        --out experiments/P0-pilot/data/results.jsonl --dry-run          # no API calls: shows the plan
    python experiments/P0-pilot/run_pilot.py --items ... --out ... --max-calls 400

Detectors (all see the passage and/or the Lean code, never the human labels):
  lint    static patterns on the projected simplex          (no LLM)
  role    source-side simplex (LLM sees only the passage) aligned to the Lean simplex
  judge   one LLM call that sees passage + Lean and says faithful / not
  compile the stored Lean verdict (``verified``), read from the items file

Hard stops: --max-calls bounds API calls; a rerun skips items already in --out.
"""
from __future__ import annotations

import argparse
import fcntl
import hashlib
import inspect
import json
import os
import sys
from pathlib import Path

from leanbench import srcsimplex
from leanbench.align import align
from leanbench.config import Settings
from leanbench.lint import lint
from leanbench.llm import OpenAICompatibleClient, Usage
from leanbench.project import project

JUDGE_PROMPT_VERSION = "judge-v1"
JUDGE_SYSTEM = (
    "You check a Lean 4 (Mathlib) formalization against a textbook passage. Decide whether the "
    "Lean statement says what the passage says: same objects, same hypotheses, same conclusion, "
    "same quantifiers. Ignore whether it compiles. Reply with JSON only: "
    '{"faithful": true|false, "confidence": 0..1, "reason": "one sentence"}'
)


def judge(client, passage: str, lean: str) -> dict:
    reply = client.complete_json([
        {"role": "system", "content": JUDGE_SYSTEM},
        {"role": "user", "content": f"Passage:\n{passage}\n\nLean:\n{lean}"},
    ])
    if not isinstance(reply, dict) or type(reply.get("faithful")) is not bool:
        raise ValueError(f"unusable judge reply: {str(reply)[:200]}")
    return {"faithful": reply["faithful"], "confidence": reply.get("confidence"),
            "reason": str(reply.get("reason", ""))[:300]}


def provider_blocked(error: Exception) -> bool:
    text = str(error).lower()
    return any(code in text for code in ("arrearage", "overdue-payment", "invalid_api_key", "http 401", "http 403"))


def complete_row(row: dict) -> bool:
    return not row.get("errors") and "role" in row and "judge" in row


class ProviderBlocked(RuntimeError):
    """A billing or authentication failure requires external recovery before resuming."""

    def __init__(self, row):
        super().__init__("Provider billing/authentication blocked; partial results were preserved")
        self.row = row


class CallBudgetExceeded(RuntimeError):
    """Stop before issuing a request that would exceed the cumulative budget."""


def run_item(client, item: dict, previous: dict | None = None) -> dict:
    if not item.get("source_text", "").strip():
        raise ValueError("Missing source_text; use a private source-inclusive export")
    code = item.get("lean_code") or ""
    row: dict = {**(previous or {}), "id": item["id"], "compiled": bool(item.get("verified"))}
    row["errors"] = dict(row.get("errors", {}))
    proj = project(code)
    row["projection"] = proj.as_dict()
    row["lint"] = lint(proj)
    if "role" not in row:
        try:
            src = srcsimplex.extract_source_simplex(client, item["source_text"], item.get("label", ""))
            al = align(src, proj)
            row["source"] = [r.__dict__ for r in src.roles]
            row["role"] = al.as_dict()
            row["errors"].pop("role", None)
        except CallBudgetExceeded as error:
            error.row = row
            raise
        except Exception as error:
            row["errors"]["role"] = str(error)[:300]
            if provider_blocked(error):
                raise ProviderBlocked(row) from error
    if "judge" not in row:
        try:
            row["judge"] = judge(client, item["source_text"], code)
            row["errors"].pop("judge", None)
        except CallBudgetExceeded as error:
            error.row = row
            raise
        except Exception as error:
            row["errors"]["judge"] = str(error)[:300]
            if provider_blocked(error):
                raise ProviderBlocked(row) from error
    return row


def digest(path: Path) -> str:
    return hashlib.sha256(path.read_bytes()).hexdigest()


def atomic_json(path: Path, value: dict) -> None:
    temporary = path.with_suffix(path.suffix + ".tmp")
    temporary.write_text(json.dumps(value, indent=2, allow_nan=False) + "\n", encoding="utf-8")
    os.replace(temporary, path)


class AuditedClient(OpenAICompatibleClient):
    """Persist reservations before HTTP requests, including transport and JSON retries."""

    def complete(self, *args, **kwargs):
        try:
            return super().complete(*args, **kwargs)
        finally:
            self.checkpoint()


def execute(args, items, settings) -> int:
    rows = []
    if args.out.exists():
        rows = [json.loads(x) for x in args.out.read_text(encoding="utf-8").splitlines() if x.strip()]
    ids = {item["id"] for item in items}
    by_id = {row["id"]: row for row in rows}
    done = {row["id"] for row in rows if complete_row(row)}
    if len(by_id) != len(rows) or not set(by_id) <= ids:
        raise ValueError("Results have duplicate or unknown item IDs")
    hashes = {
        "source": hashlib.sha256((srcsimplex.SYSTEM_PROMPT + inspect.getsource(srcsimplex.passage_prompt)).encode()).hexdigest(),
        "judge": hashlib.sha256(JUDGE_SYSTEM.encode()).hexdigest(),
    }
    identity = {"model": settings.llm.model, "base_url": settings.llm.base_url,
                "temperature": settings.llm.temperature, "items_sha256": digest(args.items),
                "prompt_hashes": hashes,
                "prompts": [srcsimplex.PROMPT_VERSION, JUDGE_PROMPT_VERSION]}
    labels = getattr(args, "labels", None)
    if labels is not None:
        identity["labels_sha256"] = digest(labels)
    meta_path = args.out.with_suffix(".meta.json")
    meta = {**identity, "request_attempts": 0, "usage": Usage().as_dict(), "status": "ready"}
    if meta_path.exists():
        meta = json.loads(meta_path.read_text(encoding="utf-8"))
        if any(meta.get(k) != v for k, v in identity.items()):
            raise ValueError("Run identity changed; use a new output path instead of mixing runs")
    elif rows:
        raise ValueError("Cannot resume results without their audit metadata")
    todo = [item for item in items if item["id"] not in done]
    missing = sum(2 - int("role" in by_id.get(i["id"], {}))
                  - int("judge" in by_id.get(i["id"], {})) for i in todo)
    print(f"items={len(items)} done={len(done)} todo={len(todo)} model={settings.llm.model} "
          f"missing_components={missing} request_attempts={meta['request_attempts']} cap={args.max_calls}")
    print("prompt_hashes:", json.dumps(hashes))
    if args.dry_run:
        return 0
    if getattr(args, "audit_only", False):
        meta["completed_items"] = len(done)
        meta["processed_items"] = len(rows)
        meta["status"] = ("provider_blocked" if any(provider_blocked(RuntimeError(message))
                          for row in rows for message in row.get("errors", {}).values())
                          else "detector_errors" if todo else "complete")
        if args.out.exists():
            meta["results_sha256"] = digest(args.out)
        atomic_json(meta_path, meta)
        print(json.dumps(meta))
        return 0 if meta["status"] == "complete" else 2
    if settings.llm.model != "qwen-max" or settings.llm.temperature != 0:
        raise ValueError("P0 is preregistered for qwen-max at temperature 0")
    client = AuditedClient.from_settings(settings.llm)
    client.usage = Usage(**meta["usage"])

    def checkpoint():
        meta["usage"] = client.usage.as_dict()
        meta["completed_items"] = len(done)
        meta["processed_items"] = len(by_id)
        atomic_json(meta_path, meta)

    def reserve():
        if meta["request_attempts"] >= args.max_calls:
            raise CallBudgetExceeded("Cumulative request cap reached; no further HTTP request was sent")
        meta["request_attempts"] += 1
        checkpoint()

    client.checkpoint = checkpoint
    client.request_guard = reserve
    meta["max_calls"] = args.max_calls
    meta["status"] = "running"
    checkpoint()
    def save_row(row):
        by_id[row["id"]] = row
        if complete_row(row):
            done.add(row["id"])
        temporary = args.out.with_suffix(args.out.suffix + ".tmp")
        with temporary.open("w", encoding="utf-8") as stream:
            for item in items:
                if item["id"] in by_id:
                    stream.write(json.dumps(by_id[item["id"]], ensure_ascii=False) + "\n")
            stream.flush()
            os.fsync(stream.fileno())
        os.replace(temporary, args.out)
        checkpoint()

    try:
        for item in todo:
            row = run_item(client, item, by_id.get(item["id"]))
            save_row(row)
            print(f"completed={len(done)}/{len(items)} requests={meta['request_attempts']}", flush=True)
        meta["status"] = "complete" if len(done) == len(items) else "detector_errors"
    except (CallBudgetExceeded, ProviderBlocked) as error:
        if getattr(error, "row", None) is not None:
            save_row(error.row)
        meta["status"] = "provider_blocked" if isinstance(error, ProviderBlocked) else "budget_exhausted"
        print(str(error), file=sys.stderr)
    except BaseException:
        meta["status"] = "interrupted"
        raise
    finally:
        if args.out.exists():
            meta["results_sha256"] = digest(args.out)
        checkpoint()
    print(json.dumps(meta))
    return 0 if meta["status"] == "complete" else 2


def main() -> int:
    ap = argparse.ArgumentParser(description=__doc__.splitlines()[0])
    ap.add_argument("--items", required=True, type=Path)
    ap.add_argument("--out", required=True, type=Path)
    ap.add_argument("--max-calls", type=int, default=400)
    ap.add_argument("--dry-run", action="store_true")
    ap.add_argument("--audit-only", action="store_true", help="Reconcile saved state without making API calls")
    ap.add_argument("--labels", type=Path, help="Frozen labels; their hash is bound to the run")
    args = ap.parse_args()
    if not 1 <= args.max_calls <= 400:
        ap.error("P0 requires a request cap between 1 and 400")
    items = [json.loads(x) for x in args.items.read_text(encoding="utf-8").splitlines() if x.strip()]
    if not items or len({item["id"] for item in items}) != len(items):
        ap.error("Input must contain unique item IDs")
    if any(not item.get("source_text", "").strip() for item in items):
        ap.error("Every item needs nonempty source_text; use a private source-inclusive export")
    if not args.dry_run and args.labels is None:
        ap.error("Freeze all blind labels first, then pass --labels")
    if args.labels is not None:
        annotation = json.loads(args.labels.read_text(encoding="utf-8"))
        key_path = args.labels.with_name("key.json")
        key = json.loads(key_path.read_text(encoding="utf-8"))["key"]
        verdicts = {"faithful", "minor_drift", "wrong", "not_formalizable"}
        if (set(key.values()) != {item["id"] for item in items}
                or set(annotation["labels"]) != set(key)
                or any(row.get("verdict") not in verdicts for row in annotation["labels"].values())
                or annotation.get("items_sha256") != digest(args.items)
                or annotation.get("key_sha256") != digest(key_path)):
            ap.error("Frozen labels, key and source-inclusive items must be complete and match")
    args.out.parent.mkdir(parents=True, exist_ok=True)
    with args.out.with_suffix(args.out.suffix + ".lock").open("w") as lock:
        fcntl.flock(lock, fcntl.LOCK_EX | fcntl.LOCK_NB)
        return execute(args, items, Settings.from_env())


if __name__ == "__main__":
    raise SystemExit(main())