#!/usr/bin/env python3
"""exp15 — pre-paid WIRE-CHECK gate (§1.13-C). FIRST paid calls (freeze committed b9272a1 already).

Confirms the three at-risk behaviors BEFORE spending on arms/judges:
  - grok-4.3 judge runs reasoning-ON (reasoning tokens > 0)   [orchestrator: confirmed reasoning-ON]
  - mistral-large-latest slug resolves + reports its version   [freeze: live-factcheck]
  - glm-5 is callable PAYG with our ZAI key (not gated to Coding Plan)  [freeze risk glm5_payg_access]

Tiny prompts (~1 sentence); per-call predicted cost printed + gated; usage/cost persisted (D30.10).
Run: eval/.venv/bin/python eval/exp15/wirecheck.py
"""
from __future__ import annotations
import json
import os
from pathlib import Path

from openai import OpenAI
from dotenv import load_dotenv

EVAL = Path(__file__).resolve().parent.parent
load_dotenv(EVAL / ".env")
FACTS = json.load(open("/home/ubuntu/books/gu-zhenren/exp15/freeze_facts.json"))["prices"]
OUT = Path("/home/ubuntu/books/gu-zhenren/exp15/wirecheck.json")
PER_CALL_CAP = 0.05  # $ predicted per-call gate

PROMPT = [{"role": "system", "content": "Ты переводчик. Выведи ТОЛЬКО перевод."},
          {"role": "user", "content": "Переведи на русский: 方源笑了，握紧拳头。"}]


def price(model, field):
    v = FACTS[model][field]
    return v[0] if isinstance(v, list) else v


def cost(model, usage):
    pin, pout = price(model, "in"), price(model, "out")
    return (usage.get("prompt_tokens", 0) * pin + usage.get("completion_tokens", 0) * pout) / 1e6


def reasoning_tokens(usage_obj):
    # xAI/OpenAI-style nested; be defensive across shapes
    d = usage_obj if isinstance(usage_obj, dict) else usage_obj.__dict__
    for k in ("reasoning_tokens",):
        if d.get(k):
            return d[k]
    ctd = d.get("completion_tokens_details")
    if ctd:
        ctd = ctd if isinstance(ctd, dict) else getattr(ctd, "__dict__", {})
        return ctd.get("reasoning_tokens", 0) or 0
    return 0


def probe(name, base, env, model, extra=None, expect_reasoning=False):
    key = os.environ.get(env)
    if not key:
        return {"name": name, "error": f"no key {env}"}
    cli = OpenAI(base_url=base, api_key=key)
    kw = dict(model=model, messages=PROMPT, max_tokens=1200)
    if extra:
        kw.update(extra)
    try:
        r = cli.chat.completions.create(**kw)
    except Exception as e:  # noqa: BLE001
        return {"name": name, "model": model, "error": f"{type(e).__name__}: {str(e)[:200]}"}
    u = r.usage
    usage = {"prompt_tokens": u.prompt_tokens, "completion_tokens": u.completion_tokens,
             "total_tokens": u.total_tokens}
    rt = reasoning_tokens(u)
    txt = (r.choices[0].message.content or "").strip()
    c = cost(model, usage)
    res = {"name": name, "model_requested": model, "model_returned": getattr(r, "model", None),
           "usage": usage, "reasoning_tokens": rt, "cost_usd": round(c, 6),
           "output_sample": txt[:120], "finish": r.choices[0].finish_reason}
    if expect_reasoning:
        res["reasoning_on_OK"] = rt > 0
    return res


def main():
    checks = [
        # grok judge: NO reasoning_effort -> default (reasoning ON). Confirm reasoning tokens > 0.
        dict(name="grok-judge", base="https://api.x.ai/v1", env="XAI_API_KEY",
             model="grok-4.3", expect_reasoning=True),
        # mistral judge/arm: confirm slug resolves + version returned
        dict(name="mistral-large", base="https://api.mistral.ai/v1", env="MISTRAL_API_KEY",
             model="mistral-large-latest"),
        # glm-5 editor: confirm PAYG callable (thinking off for a clean quick probe)
        dict(name="glm-5-editor", base="https://api.z.ai/api/paas/v4", env="ZAI_API_KEY",
             model="glm-5", extra={"extra_body": {"thinking": {"type": "disabled"}}}),
    ]
    results, total = [], 0.0
    for ch in checks:
        # per-call predicted cost gate (worst case max_tokens fully used as output)
        pred = price(ch["model"], "out") * 1200 / 1e6 + price(ch["model"], "in") * 60 / 1e6
        print(f"[gate] {ch['name']} ({ch['model']}) predicted <= ${pred:.4f}  {'OK' if pred<=PER_CALL_CAP else 'BLOCKED'}")
        if pred > PER_CALL_CAP:
            results.append({"name": ch["name"], "error": f"per-call gate blocked (pred ${pred:.4f} > ${PER_CALL_CAP})"})
            continue
        res = probe(**ch)
        results.append(res)
        total += res.get("cost_usd", 0) or 0
        if "error" in res:
            print(f"  ERROR: {res['error']}")
        else:
            extra = f" reasoning_tokens={res['reasoning_tokens']}" if res.get("reasoning_tokens") else ""
            print(f"  OK model_returned={res['model_returned']} cost=${res['cost_usd']:.5f}{extra}")
            print(f"     out: {res['output_sample']!r}")

    OUT.write_text(json.dumps({"total_cost_usd": round(total, 6), "results": results}, ensure_ascii=False, indent=2),
                   encoding="utf-8")
    print(f"\ntotal wire-check spend: ${total:.5f}  -> {OUT}")

    # self-review asserts
    by = {r["name"]: r for r in results}
    grok = by.get("grok-judge", {})
    if "error" not in grok:
        assert grok.get("reasoning_on_OK"), f"grok reasoning NOT on (reasoning_tokens={grok.get('reasoning_tokens')})"
        print("OK: grok-4.3 judge reasoning-ON confirmed.")
    if "error" not in by.get("glm-5-editor", {}):
        print("OK: glm-5 PAYG callable with ZAI key.")
    if "error" not in by.get("mistral-large", {}):
        print(f"OK: mistral-large-latest -> {by['mistral-large'].get('model_returned')}")


if __name__ == "__main__":
    main()