#!/usr/bin/env python3
"""exp13 — бейк-офф ПЕРЕВОДЧИКОВ (платно, кап $5). Пре-регистрация §1.2 заморожена.

Faithful draft-стадия: swap ТОЛЬКО модели переводчика. Layout повторяет прод draft-стадию
(chunkrun.go: translatorInjection = renderGlossaryBlock; render.go MessagesWithInjection):
  system(translator.md reflow-edition, brief-filled) → system(инъекция «src → dst») → user(перевод).
Инъекция и исходник — из приёмочного store (V1-сид), ИДЕНТИЧНЫ по чанку у всех армов.
translator.md строится в коде с УДАЛЕНИЕМ строки 9 «Сохраняй разбивку…» (ратифиц. D30.2;
пакет бэкенда ещё не залендён — прод-файл не трогаю).

T1 deepseek-v4-flash — из store ($0). T6 qwen — только с DashScope-ключом (env QWEN_API_KEY|DASHSCOPE_API_KEY).
Запуск: eval/.venv/bin/python eval/exp13_translate.py [--dry]
Выход: diag/arms13/<ARM>/<ch>_<ci>.txt + diag/arms13/results.json + diag/arms13/costs.jsonl
"""
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.parent))  # script now under eval/expNN/ — reach up to parent eval/ for the refusal_bench oracle
from refusal_bench import load_env_file
load_env_file()

DIAG = Path("/home/ubuntu/books/gu-zhenren/diag")
OUT = DIAG / "arms13"
TRANSLATOR_MD = Path("/home/ubuntu/projects/textmachine/backend/prompts/translator.md")
CAP_USD = 5.0
TARGETS = [(5, 0), (5, 1), (7, 0), (7, 1), (14, 0), (14, 1)]

# brief-vars из acceptance/book.yaml (D18)
VARS = {"source_lang": "zh", "target_lang": "ru", "genre": "вебновелла",
        "audience": "взрослые читатели вебновелл", "title": "蛊真人",
        "honorifics": "keep", "transcription": "palladius", "venuti": "0.6",
        "footnotes": "minimal"}

# строка D30.2-reflow: удаляем ровно её (ратифицированный снос «Сохраняй разбивку»)
REFLOW_DROP = "Сохраняй разбивку на абзацы и прямую речь как в оригинале."

PRICES = {  # $/1M in/out (models.yaml, prices_checked 2026-07-10)
    "deepseek-v4-flash": (0.14, 0.28), "deepseek-v4-pro": (0.435, 0.87),
    "mistral-large-2512": (0.5, 1.5), "glm-5": (1.0, 3.2), "grok-4.3": (1.25, 2.50),
}

ARMS = {
    # T1 — из store, генерации нет (см. main); прайс для справки
    "T2": dict(model="mistral-large-2512", base="https://api.mistral.ai/v1",
               keyenv="MISTRAL_API_KEY", extra={}, max_tokens=8000),
    "T3": dict(model="glm-5", base="https://api.z.ai/api/paas/v4",
               keyenv="ZAI_API_KEY", extra={"thinking": {"type": "disabled"}}, max_tokens=8000),
    "T4": dict(model="deepseek-v4-pro", base="https://api.deepseek.com/v1",
               keyenv="DEEPSEEK_API_KEY", extra={}, max_tokens=8000),   # thinking ON = default
    "T5": dict(model="grok-4.3", base="https://api.x.ai/v1",
               keyenv="XAI_API_KEY", extra={}, max_tokens=8000),        # reasoning ON = default low
}
# T6 qwen — гейт на ключ; DashScope OpenAI-совместимый слой
QWEN = dict(model="qwen-flash", base="https://dashscope-intl.aliyuncs.com/compatible-mode/v1",
            keyenv="DASHSCOPE_API_KEY", extra={}, max_tokens=8000)


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


def translator_system() -> str:
    raw = TRANSLATOR_MD.read_text(encoding="utf-8")
    sys_part = raw.split("---USER---")[0].rstrip()
    lines = [ln for ln in sys_part.split("\n") if REFLOW_DROP not in ln]
    dropped = sys_part.count(REFLOW_DROP)
    if dropped != 1:
        raise SystemExit(f"reflow-guard: ожидал 1 вхождение строки D30.2, нашёл {dropped} — translator.md изменился, сверься с оркестратором")
    return render("\n".join(lines))


def build_messages(source: str, inj: str) -> list[dict]:
    msgs = [{"role": "system", "content": translator_system()}]
    if inj.strip():
        msgs.append({"role": "system", "content": inj})
    msgs.append({"role": "user", "content": f"Переведи следующий фрагмент:\n\n{source}"})
    return msgs


def call(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.3, "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
    reason = usage.get("reasoning_tokens", 0) or 0  # grok additive шлёт отдельно; subset=0
    pin, pout = PRICES.get(model, (0, 0))
    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"))
    OUT.mkdir(parents=True, exist_ok=True)
    costs_fp = OUT / "costs.jsonl"
    results, spent = [], 0.0

    arms = dict(ARMS)
    qkey = os.environ.get("DASHSCOPE_API_KEY") or os.environ.get("QWEN_API_KEY")
    if qkey:
        os.environ["DASHSCOPE_API_KEY"] = qkey
        arms["T6"] = QWEN
        print("T6 qwen: DashScope-ключ найден — включаю")
    else:
        print("T6 qwen: ключа нет — пропускаю (не блокируюсь)")

    # T1 — из store (deepseek-v4-flash draft), $0
    t1 = OUT / "T1"; t1.mkdir(exist_ok=True)
    for ch, ci in TARGETS:
        (t1 / f"{ch}_{ci}.txt").write_text(chunks[f"{ch}/{ci}"]["draft"], encoding="utf-8")
    print("T1 deepseek-v4-flash: 6 чанков из store ($0)")

    for arm, spec in arms.items():
        adir = OUT / arm; adir.mkdir(exist_ok=True)
        for ch, ci in TARGETS:
            key = f"{ch}/{ci}"
            src = chunks[key]["source"]
            inj = blocks[key]["bilingual_glossary"]
            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(src, inj)
            if dry:
                approx_in = sum(len(m["content"]) for m in msgs) // 3
                print(f"  [{arm} {key}] model={spec['model']} msgs={len(msgs)} ~in={approx_in}tok src_chars={len(src)}")
                continue
            if spent > CAP_USD * 0.8:
                print(f"!! near cap (${spent:.3f}) — stop"); break
            t0 = time.time()
            text, err, usage = call(spec, msgs)
            c = cost_of(spec["model"], usage)
            spent += c
            rec = {"arm": arm, "model": spec["model"], "chunk": key, "cost": round(c, 6),
                   "usage": {k: usage.get(k) for k in ("prompt_tokens", "completion_tokens", "reasoning_tokens", "total_tokens", "finish_reason")}}
            with costs_fp.open("a", encoding="utf-8") as f:
                f.write(json.dumps(rec, ensure_ascii=False) + "\n")
            if not text:
                print(f"  [{arm} {key}] ERROR: {err} (${spent:.3f})")
                results.append({**rec, "error": err}); continue
            outp.write_text(text, encoding="utf-8")
            results.append({**rec, "chars": len(text)})
            print(f"  [{arm} {key}] {time.time()-t0:.0f}s {len(text)}ch ${c:.4f} cum=${spent:.3f} finish={usage.get('finish_reason')}")

    if not dry:
        (OUT / "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()