#!/usr/bin/env python3
"""B3(+B2): пробы консилиума по пре-регистрации research/24 §0. Прямой API, не движок.

Пассы: ds1..ds3 (deepseek flash, effort low, само-согласованность) · dsw (словесная шкала) ·
glm1 (glm-5, thinking off) · mis1 (mistral-large-2512) · gro1 (grok-4.3, опция --grok).
Вход — engine-faithful RenderBatch-блоки из bank-stop coldrun-a (серии со-батчены, как §1 движка).
Сырьё — durable ~/books/gu-zhenren/bank-arbitration/. Деньги считаются из usage × цены models.yaml.
"""
from __future__ import annotations
import argparse
import collections
import json
import os
import sys
import time
from pathlib import Path

from dotenv import load_dotenv
from openai import OpenAI

REPO = Path("/home/ubuntu/projects/textmachine")
HERE = Path(__file__).resolve().parent
RAW = Path.home() / "books" / "gu-zhenren" / "bank-arbitration"
RAW.mkdir(exist_ok=True)
load_dotenv(REPO / "eval" / ".env")  # ключи грузит скрипт; значения наружу не печатаются

sys.path.insert(0, str(HERE.parent / "bank_autonomy"))
from parse_common import parse_bankstop  # noqa: E402

# Цены — backend/configs/models.yaml (prices_checked 2026-07-10): input/cached/output за 1M.
PRICES = {
    "deepseek-v4-flash": (0.14, 0.0028, 0.28),
    "glm-5": (1.0, 0.2, 3.2),
    "mistral-large-2512": (0.5, 0.5, 1.5),
    "grok-4.3": (1.25, 1.25, 2.50),
}
PROVIDERS = {
    "deepseek-v4-flash": ("https://api.deepseek.com/v1", "DEEPSEEK_API_KEY"),
    "glm-5": ("https://api.z.ai/api/paas/v4", "ZAI_API_KEY"),
    "mistral-large-2512": ("https://api.mistral.ai/v1", "MISTRAL_API_KEY"),
    "grok-4.3": ("https://api.x.ai/v1", "XAI_API_KEY"),
}

TERM_PROMPT = REPO / "backend/prompts/zh-ru/terminologist.md"
USER_SEP = "\n---USER---\n"
BATCH_RUNES = 6000  # terminologyDefaultBatchRunes

CONF_NUM = ("\n\nДОПОЛНЕНИЕ К ФОРМАТУ (только для этого прогона): добавь ТРЕТЬЕ поле через "
            "табуляцию — твою уверенность в выбранном переводе ЧИСЛОМ от 0 до 100 "
            "(0 — наугад, 100 — абсолютно уверен). Пример: 师父\tнаставник\t85")
CONF_WORD = ("\n\nДОПОЛНЕНИЕ К ФОРМАТУ (только для этого прогона): добавь ТРЕТЬЕ поле через "
             "табуляцию — твою уверенность в выбранном переводе РОВНО ОДНИМ из выражений: "
             "точно | скорее да | не уверен | не знаю. Пример: 师父\tнаставник\tточно")

# Экстра-поверхности для кластерной метрики (все живут в BANK-FULL; голд-dst у них нет,
# скорим только связность конвенции внутри семьи).
EXTRAS = [
    "二转", "四转", "九转",               # достройка 转-серии
    "甲等资质", "乙等资质", "丙等资质",     # серия 3
    "初阶", "中阶", "高阶",               # серия 4
    "古月师", "古月族长", "古月一族", "古月陈博", "古月藻榭", "古月山寨",  # 古月-семья (разнодлинная)
    "蛊室", "本命蛊", "希望蛊", "炼蛊", "一转蛊师", "三转蛊师",           # 蛊-семья
]

SERIES = {
    "zhuan": ["一转", "二转", "三转", "四转", "五转", "六转", "九转"],  # 二转 из BANK-FULL, добран в EXTRAS
    "deng": ["甲等", "乙等", "丙等", "丁等"],
    "dengzizhi": ["甲等资质", "乙等资质", "丙等资质"],
    "jie": ["初阶", "中阶", "高阶"],
}


def render_terminologist(text_block: str) -> tuple[str, str]:
    raw = TERM_PROMPT.read_text(encoding="utf-8")
    out, rest = [], raw
    while True:
        i = rest.find("<!--")
        if i < 0:
            out.append(rest)
            break
        out.append(rest[:i])
        rest = rest[i + 4:]
        j = rest.find("-->")
        rest = rest[j + 3:]
    canon = "".join(out)
    head, user = canon.split(USER_SEP, 1)
    vals = {"source_lang": "zh", "target_lang": "ru", "genre": "вебновелла",
            "audience": "взрослые читатели вебновелл", "title": "蛊真人", "venuti": "0.60",
            "honorifics": "keep", "transcription": "palladius", "footnotes": "minimal",
            "text": text_block}
    system = head.strip()
    for k, v in vals.items():
        system = system.replace("{{" + k + "}}", v)
        user = user.replace("{{" + k + "}}", v)
    return system, user.strip()


def block_of(c: dict) -> str:
    b = [f"### {c['src']}",
         f"key: {c['src']}\ntype: {c['type'] or 'term'}\norigin: {c['origin']}\nfreq: {c['freq']}\nsince_ch: 0"]
    if c.get("evidence"):
        b.append(f"evidence: {', '.join(c['evidence'])}")
    if c.get("drafts"):
        b.append("drafts: " + " | ".join(f"{d} ×{n}" for d, n in c["drafts"]))
    for k in c.get("ctx", []):
        b.append(f"ctx: {k}")
    return "\n".join(b)


def build_roster() -> list[dict]:
    bs = parse_bankstop()
    gold = [json.loads(l) for l in open(HERE / "gold" / "gold.jsonl", encoding="utf-8")]
    srcs = [r["bank_src"] for r in gold if r.get("in_bank")]
    for e in EXTRAS:
        if e not in srcs:
            srcs.append(e)
    roster = []
    missing = []
    for s in srcs:
        if s not in bs:
            missing.append(s)
            continue
        r = bs[s]
        roster.append({"src": s, "type": r.get("type", "term"), "origin": r.get("origin", ""),
                       "freq": r.get("freq", 0), "drafts": r.get("drafts", []),
                       "ctx": r.get("ctx", []), "evidence": r.get("evidence", [])})
    if missing:
        print("NOT in bank-stop (пропущены):", ", ".join(missing))
    return roster


def batches_of(roster: list[dict]) -> list[list[dict]]:
    """Реплика Batch(): ключ-сортировка → серии contiguous → пак ≤6000 рун → батчи по freq desc."""
    sid = {}
    for i, members in enumerate(SERIES.values(), 1):
        for m in members:
            sid[m] = i
    ordered = sorted(roster, key=lambda c: c["src"])
    # series contiguous, anchored at first member position
    out_order, emitted = [], set()
    by_sid = collections.defaultdict(list)
    for c in ordered:
        if sid.get(c["src"]):
            by_sid[sid[c["src"]]].append(c)
    for c in ordered:
        s = sid.get(c["src"])
        if not s:
            out_order.append(c)
        elif s not in emitted:
            out_order.extend(by_sid[s])
            emitted.add(s)
    batches, cur, size = [], [], 0
    i = 0
    while i < len(out_order):
        j = i + 1
        if sid.get(out_order[i]["src"]):
            s = sid[out_order[i]["src"]]
            while j < len(out_order) and sid.get(out_order[j]["src"]) == s:
                j += 1
        unit = out_order[i:j]
        n = sum(len(block_of(c)) + 2 for c in unit)
        if cur and size + n > BATCH_RUNES:
            batches.append(cur)
            cur, size = [], 0
        cur.extend(unit)
        size += n
        i = j
    if cur:
        batches.append(cur)
    batches.sort(key=lambda b: -sum(c["freq"] for c in b))
    return batches


def call(client: OpenAI, model: str, system: str, user: str, tag: str,
         effort: str | None, thinking_off: bool, cap: int = 16000) -> dict:
    msgs = [{"role": "system", "content": system}, {"role": "user", "content": user}]
    kw = dict(model=model, messages=msgs, max_tokens=cap, temperature=0)
    extra = {}
    if effort:
        extra["reasoning_effort"] = effort
    if thinking_off:
        extra["thinking"] = {"type": "disabled"}
    if extra:
        kw["extra_body"] = extra
    t0 = time.time()
    r = client.chat.completions.create(**kw)
    ch = r.choices[0]
    content = ch.message.content or ""
    reasoning = getattr(ch.message, "reasoning_content", None) or ""
    u = r.usage
    pt = getattr(u, "prompt_tokens", 0) or 0
    ct = getattr(u, "completion_tokens", 0) or 0
    cached = getattr(getattr(u, "prompt_tokens_details", None), "cached_tokens", 0) or 0
    pin, pcache, pout = PRICES[model]
    cost = (pt - cached) / 1e6 * pin + cached / 1e6 * pcache + ct / 1e6 * pout
    rec = dict(tag=tag, model=model, model_returned=r.model, effort=effort or ("thinking-off" if thinking_off else "default"),
               ts=time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()), cap=cap,
               finish=ch.finish_reason, prompt_tokens=pt, cached_tokens=cached,
               completion_tokens=ct, cost_usd=round(cost, 6), latency_s=round(time.time() - t0, 1),
               content=content, reasoning_chars=len(reasoning), system=system, user=user)
    (RAW / f"{tag}.json").write_text(json.dumps(rec, ensure_ascii=False, indent=1), encoding="utf-8")
    print(f"[{tag}] {model} finish={ch.finish_reason} in={pt}(c{cached}) out={ct} ${cost:.6f} {rec['latency_s']}s content={len(content)}c")
    return rec


def slug_check(model: str) -> None:
    base, keyenv = PROVIDERS[model]
    key = os.environ.get(keyenv)
    if not key:
        print(f"SKIP {model}: нет {keyenv}")
        raise KeyError(keyenv)
    cl = OpenAI(api_key=key, base_url=base, timeout=120)
    ids = [m.id for m in cl.models.list().data]
    ok = model in ids
    print(f"/models {model}: {'OK' if ok else 'ОТСУТСТВУЕТ! listing=' + ','.join(ids[:20])}")
    if not ok:
        raise RuntimeError(f"слаг {model} не листится")


def client_for(model: str) -> OpenAI:
    base, keyenv = PROVIDERS[model]
    return OpenAI(api_key=os.environ[keyenv], base_url=base, timeout=600)


PASSES = {
    # tag: (model, effort, thinking_off, conf_instruction)
    "ds1": ("deepseek-v4-flash", "low", False, CONF_NUM),
    "ds2": ("deepseek-v4-flash", "low", False, CONF_NUM),
    "ds3": ("deepseek-v4-flash", "low", False, CONF_NUM),
    "dsw": ("deepseek-v4-flash", "low", False, CONF_WORD),
    "glm1": ("glm-5", None, True, CONF_NUM),
    "mis1": ("mistral-large-2512", None, False, CONF_NUM),
    "gro1": ("grok-4.3", None, False, CONF_NUM),
}


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--passes", default="ds1,ds2,ds3,dsw,glm1,mis1")
    ap.add_argument("--dry", action="store_true", help="напечатать батчи и смету, без вызовов")
    a = ap.parse_args()

    roster = build_roster()
    batches = batches_of(roster)
    print(f"roster={len(roster)} терм., batches={len(batches)}: " +
          " ".join(f"#{i}:{len(b)}терм/{sum(len(block_of(c))+2 for c in b)}рун" for i, b in enumerate(batches)))
    if a.dry:
        for i, b in enumerate(batches):
            print(f"--- batch {i}: {' '.join(c['src'] for c in b)}")
        return

    want = [p.strip() for p in a.passes.split(",") if p.strip()]
    for model in sorted({PASSES[p][0] for p in want}):
        slug_check(model)

    total = 0.0
    for p in want:
        model, effort, toff, conf = PASSES[p]
        cl = client_for(model)
        for i, b in enumerate(batches):
            system, user = render_terminologist("\n\n".join(block_of(c) for c in b))
            rec = call(cl, model, system, user + conf, f"{p}-b{i}", effort, toff)
            total += rec["cost_usd"]
            time.sleep(1.0)
    print(f"TOTAL cost this run: ${total:.6f}")


if __name__ == "__main__":
    main()