#!/usr/bin/env python3
"""exp12 — генерация редакторских армов A1′,A2,A3,A4,A5,A6 (платно, кап $5).

Черновики (A0) и финал glm-5-моно (A1) — из store, $0 (не тут). Здесь ре-редактируем
ТЕ ЖЕ черновики разными редакторами/режимами. Промпты ЗАМОРОЖЕНЫ (пре-регистрация §1.2).

Faithful render: system(editor.md) → system(инъекция-блок) → user(черновик[+исходник]).
Layout повторяет render.go MessagesWithInjection.

Запуск: eval/.venv/bin/python eval/exp12_arms.py [--dry]
Выход: diag/arms/<ARM>/<ch>_<chunk>.txt + diag/arms/results.json
"""
from __future__ import annotations
import json, os, sys, time, urllib.request, urllib.error
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parent))
from refusal_bench import load_env_file
load_env_file()

DIAG = Path("/home/ubuntu/books/gu-zhenren/diag")
ARMS = DIAG / "arms"
EDITOR_MD = Path("/home/ubuntu/projects/textmachine/backend/prompts/editor.md")
CAP_USD = 5.0

# --- render vars (book.yaml) ---
VARS = {"genre": "вебновелла", "audience": "взрослые читатели вебновелл",
        "title": "蛊真人", "source_lang": "zh", "honorifics": "keep", "venuti": "0.60"}

PRICES = {  # per 1M tokens (models.yaml, prices_checked 2026-07-10)
    "glm-5": (1.0, 3.2), "grok-4.3": (1.25, 2.50),
    "gemini-3.1-pro-preview": (2.0, 12.0),
}

# --- frozen arm specs ---
BILINGUAL_LINE = ("Тебе также дан ИСХОДНЫЙ ТЕКСТ: сверяйся с ним по смыслу и исправляй "
                  "искажения черновика, сохраняя полноту.")
ARMS_SPEC = {
    "A1p": dict(model="glm-5", base="https://api.z.ai/api/paas/v4", keyenv="ZAI_API_KEY",
                mode="mono", constraints=True, extra={"thinking": {"type": "disabled"}}, max_tokens=8000),
    "A2":  dict(model="glm-5", base="https://api.z.ai/api/paas/v4", keyenv="ZAI_API_KEY",
                mode="biling", constraints=True, extra={"thinking": {"type": "disabled"}}, max_tokens=8000),
    "A3":  dict(model="grok-4.3", base="https://api.x.ai/v1", keyenv="XAI_API_KEY",
                mode="mono", constraints=True, extra={"reasoning_effort": "none"}, max_tokens=8000),
    "A4":  dict(model="grok-4.3", base="https://api.x.ai/v1", keyenv="XAI_API_KEY",
                mode="biling", constraints=True, extra={"reasoning_effort": "none"}, max_tokens=8000),
    "A5":  dict(model="gemini-3.1-pro-preview",
                base="https://generativelanguage.googleapis.com/v1beta/openai", keyenv="GEMINI_API_KEY",
                mode="biling", constraints=True, extra={}, max_tokens=16000),
    "A6":  dict(model="grok-4.3", base="https://api.x.ai/v1", keyenv="XAI_API_KEY",
                mode="mono", constraints=False, extra={"reasoning_effort": "none"}, max_tokens=8000),
}


def render(tpl: str) -> str:
    for k, v in VARS.items():
        tpl = tpl.replace("{{%s}}" % k, v)
    return tpl


def editor_system() -> str:
    raw = EDITOR_MD.read_text(encoding="utf-8")
    sys_part = raw.split("---USER---")[0].rstrip()
    return render(sys_part)


def build_messages(spec: dict, source: str, draft: str, blk: dict) -> list[dict]:
    sys1 = editor_system()
    if spec["mode"] == "biling":
        sys1 = sys1 + "\n" + BILINGUAL_LINE
    msgs = [{"role": "system", "content": sys1}]
    if spec["constraints"]:
        inj = blk["bilingual_glossary"] if spec["mode"] == "biling" else blk["editor_constraint"]
        if inj.strip():
            msgs.append({"role": "system", "content": inj})
    if spec["mode"] == "biling":
        user = f"ИСХОДНЫЙ ТЕКСТ (zh):\n\n{source}\n\n---\n\nЧерновик перевода для редактуры:\n\n{draft}"
    else:
        user = f"Черновик перевода для редактуры:\n\n{draft}"
    msgs.append({"role": "user", "content": user})
    return msgs


def call_messages(spec: dict, messages: list[dict], timeout: int = 300):
    key = os.environ.get(spec["keyenv"], "")
    if not key:
        return None, f"no key {spec['keyenv']}", {}
    body = {"model": spec["model"], "messages": messages,
            "temperature": 0.4, "max_tokens": spec["max_tokens"], **spec["extra"]}
    req = urllib.request.Request(spec["base"].rstrip("/") + "/chat/completions",
        data=json.dumps(body).encode(),
        headers={"Content-Type": "application/json", "Authorization": f"Bearer {key}"})
    try:
        with urllib.request.urlopen(req, timeout=timeout) as r:
            data = json.loads(r.read())
    except urllib.error.HTTPError as e:
        return None, f"HTTP {e.code}: {e.read()[:200].decode('utf-8','replace')}", {}
    except Exception as e:
        return None, f"ERR {str(e)[:160]}", {}
    ch = data["choices"][0]
    text = (ch.get("message", {}) or {}).get("content") or ""
    usage = data.get("usage", {}) or {}
    usage["finish_reason"] = ch.get("finish_reason", "")
    return text.strip(), (None if text.strip() else f"empty|finish={ch.get('finish_reason')}"), usage


def cost_of(model: str, usage: dict) -> float:
    inp = usage.get("prompt_tokens", 0) or 0
    comp = usage.get("completion_tokens", 0) or 0
    total = usage.get("total_tokens", 0) or 0
    reason = usage.get("reasoning_tokens", 0) or 0
    pin, pout = PRICES.get(model, (0, 0))
    if model == "gemini-3.1-pro-preview":  # additive_total: thinking only in total_tokens
        reason = max(0, total - inp - comp)
    return (inp * pin + (comp + reason) * pout) / 1_000_000


def main() -> None:
    dry = "--dry" in sys.argv
    chunks = json.loads((DIAG / "chunks.json").read_text(encoding="utf-8"))
    blocks = json.loads((DIAG / "blocks.json").read_text(encoding="utf-8"))
    targets = [(5, 0), (5, 1), (7, 0), (7, 1), (14, 0), (14, 1)]

    ARMS.mkdir(parents=True, exist_ok=True)
    results = []
    spent = 0.0
    for arm, spec in ARMS_SPEC.items():
        adir = ARMS / arm
        adir.mkdir(exist_ok=True)
        for ch, ci in targets:
            key = f"{ch}/{ci}"
            src = chunks[key]["source"]
            draft = chunks[key]["draft"]
            blk = blocks[key]
            outp = adir / f"{ch}_{ci}.txt"
            if outp.exists() and outp.read_text(encoding="utf-8").strip():
                print(f"  [{arm} {key}] cached"); continue
            msgs = build_messages(spec, src, draft, blk)
            if dry:
                approx_in = sum(len(m["content"]) for m in msgs) // 3
                print(f"  [{arm} {key}] mode={spec['mode']} constr={spec['constraints']} "
                      f"msgs={len(msgs)} ~in={approx_in}tok draft_chars={len(draft)}")
                continue
            if spent > CAP_USD * 0.8:
                print(f"!! near cap (${spent:.3f}) — stop"); break
            t0 = time.time()
            text, err, usage = call_messages(spec, msgs)
            c = cost_of(spec["model"], usage)
            spent += c
            if not text:
                print(f"  [{arm} {key}] ERROR: {err} (${spent:.3f})")
                results.append({"arm": arm, "chunk": key, "error": err, "cost": c})
                continue
            outp.write_text(text, encoding="utf-8")
            results.append({"arm": arm, "chunk": key, "chars": len(text),
                            "finish": usage.get("finish_reason"), "usage": usage, "cost": round(c, 5)})
            print(f"  [{arm} {key}] {time.time()-t0:.0f}s {len(text)}ch ${c:.4f} cum=${spent:.3f} "
                  f"finish={usage.get('finish_reason')}")
    if not dry:
        (ARMS / "results.json").write_text(json.dumps(results, ensure_ascii=False, indent=1), encoding="utf-8")
        print(f"\nTOTAL arm spend: ${spent:.4f} (cap ${CAP_USD})")


if __name__ == "__main__":
    main()