#!/usr/bin/env python3
"""exp13 — судейская панель (пре-рег §1.4). Span-цитирующая, reasoning-ON, оси РАЗДЕЛЬНО,
кросс-семейная (ни один арм — не из семей судей), ≥3 повтора со свапом, parse-чек полноты.

Панель: J1 gemini-3.1-pro-preview (mandatory thinking), J2 gpt-5-mini (reasoning_effort=medium),
J3 kimi-k2.6 (reasoning-ON, temp=1). Каждый вызов = ранжирование ВСЕХ армов по ОДНОЙ оси на
одном чанке; JSON {ranking, evidence_spans}. Слепые метки V1..VN, свап на (чанк×ось×повтор).

Параллель внутри судьи; хард-чек спенда между судьями (кап $3.5 судьям). Персист usage/cost.
Запуск: eval/.venv/bin/python eval/exp13_judges.py [--repeats 3]
Выход: diag/arms13/judges/raw/*.txt, judges/verdicts.jsonl, judges/_JUDGE_KEY.json, judge_costs.jsonl
"""
from __future__ import annotations
import json, os, re, sys, time, random, urllib.request, urllib.error
from concurrent.futures import ThreadPoolExecutor
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 / "arms13"
JOUT = ARMS / "judges"
TARGETS = [(5, 0), (5, 1), (7, 0), (7, 1), (14, 0), (14, 1)]
REPEATS = 3
JUDGE_CAP = 3.5
PRICES = {"gemini-3.1-pro-preview": (2.0, 12.0), "gpt-5-mini": (0.25, 2.0), "kimi-k2.6": (0.95, 4.0)}

AXES = {
    "fidelity": "ВЕРНОСТЬ исходнику: точность смысла; отсутствие искажений, инверсий смысла, пропусков, отсебятины и добавлений. Стиль игнорируй.",
    "style": "СТИЛЬ русской художественной прозы: живость и читаемость, отсутствие канцелярита и калек с китайского, естественный синтаксис и ритм (школа Норы Галь). Верность НЕ оценивай.",
}

# judge_type: gemini | openai | kimi — разная wire-форма (провайдер-квирки models.yaml)
JUDGES = {
    "J1_gemini": dict(model="gemini-3.1-pro-preview", base="https://generativelanguage.googleapis.com/v1beta/openai",
                      keyenv="GEMINI_API_KEY", jtype="gemini"),
    "J2_gpt5mini": dict(model="gpt-5-mini", base="https://api.openai.com/v1",
                        keyenv="OPENAI_API_KEY", jtype="openai"),
    "J3_kimi": dict(model="kimi-k2.6", base="https://api.moonshot.ai/v1",
                    keyenv="KIMI_API_KEY", jtype="kimi"),
}


def arms_present() -> list[str]:
    return sorted([p.name for p in ARMS.iterdir() if p.is_dir() and p.name.startswith("T")])


def arm_text(arm: str, ch: int, ci: int) -> str:
    p = ARMS / arm / f"{ch}_{ci}.txt"
    return p.read_text(encoding="utf-8").strip() if p.exists() else ""


def build_labels(arms: list[str], ch: int, ci: int, axis: str, rep: int) -> dict:
    order = arms[:]
    random.Random(f"exp13-judge-{ch}-{ci}-{axis}-{rep}").shuffle(order)
    return {f"V{i+1}": arm for i, arm in enumerate(order)}


def build_messages(source: str, labels: dict, ch: int, ci: int, axis: str) -> list[dict]:
    n = len(labels)
    sys_p = ("Ты — строгий литературный судья художественного перевода китайской вебновеллы "
             "(сянься/культиваторка) на русский язык. Тебе дан ИСХОДНИК (zh) и {n} переводов под "
             "слепыми метками. Оцени и РАНЖИРУЙ переводы ТОЛЬКО по одной оси:\n\n{axis}\n\n"
             "Верни СТРОГО валидный JSON без текста вокруг:\n"
             '{{"ranking": ["метки от ЛУЧШЕЙ к ХУДШЕЙ, все {n}"], '
             '"evidence_spans": [{{"label": "Vx", "quote": "цитата ≤15 слов из перевода или исходника", '
             '"note": "чем эта деталь хороша/плоха по оси"}}]}}\n'
             "Каждой метке — минимум один span с конкретной цитатой (не пересказ). "
             "ranking должен содержать ровно метки V1..V{n}, без пропусков и повторов.").format(n=n, axis=AXES[axis])
    body = [f"ИСХОДНИК (zh):\n\n{source}\n"]
    for lab in [f"V{i+1}" for i in range(n)]:
        body.append(f"=== {lab} ===\n{arm_text(labels[lab], ch, ci)}\n")
    body.append(f"Ранжируй {n} переводов по оси «{axis}». Только JSON.")
    return [{"role": "system", "content": sys_p}, {"role": "user", "content": "\n".join(body)}]


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}
    jt = spec["jtype"]
    if jt == "gemini":
        body.update(max_tokens=8000, temperature=0.3)         # mandatory thinking (не шлём reasoning)
    elif jt == "openai":
        body.update(max_completion_tokens=8000, reasoning_effort="medium")  # gpt-5: no temp, max_completion_tokens
    elif jt == "kimi":
        body.update(max_tokens=16000, temperature=1)          # kimi: temp force 1, floor 16000
    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
    pin, pout = PRICES.get(model, (0, 0))
    if model == "gemini-3.1-pro-preview":           # additive_total: thinking только в total
        reason = max(0, total - inp - comp)
        return (inp * pin + (comp + reason) * pout) / 1_000_000
    return (inp * pin + comp * pout) / 1_000_000    # subset: comp уже включает reasoning


def parse_ranking(text: str, labels: list[str]):
    """Достаём ranking робастно; parse-чек полноты set(ranking)==set(labels).
    Судьи (gemini) ломают полный JSON неэкранированными кавычками в ru-цитатах evidence_spans —
    ranking-массив при этом чистый, извлекаем его напрямую (fallback). Spans — best-effort."""
    t = text.strip()
    if "```" in t:
        m = re.search(r"```(?:json)?\s*(.*?)```", t, re.S)
        if m:
            t = m.group(1).strip()
    obj = None
    mo = re.search(r"\{.*\}", t, re.S)
    if mo:
        for cand in (mo.group(0), mo.group(0).replace("\n", " ")):
            try:
                obj = json.loads(cand)
                break
            except Exception:
                obj = None
    rank = obj.get("ranking") if isinstance(obj, dict) else None
    if not isinstance(rank, list):
        rm = re.search(r'"ranking"\s*:\s*(\[[^\]]*\])', t)
        if rm:
            try:
                rank = json.loads(rm.group(1))
            except Exception:
                rank = None
    if not isinstance(rank, list) or set(rank) != set(labels):
        return None, f"incomplete:{rank}", obj
    return rank, None, obj


def run_judge(jname: str, spec: dict, arms: list[str], chunks: dict, keymap: dict, spent_box: list):
    tasks = [(ch, ci, axis, rep) for ch, ci in TARGETS for axis in AXES for rep in range(REPEATS)]
    raw_dir = JOUT / "raw"; raw_dir.mkdir(parents=True, exist_ok=True)
    verdicts, jcosts = [], []

    def one(t):
        ch, ci, axis, rep = t
        labels = keymap[f"{ch}/{ci}"][axis][str(rep)]
        src = chunks[f"{ch}/{ci}"]["source"]
        msgs = build_messages(src, labels, ch, ci, axis)
        text, err, usage = call(spec, msgs)
        c = cost_of(spec["model"], usage)
        tag = f"{jname}_{ch}_{ci}_{axis}_r{rep}"
        rec = {"judge": jname, "model": spec["model"], "chunk": f"{ch}/{ci}", "axis": axis,
               "rep": rep, "cost": round(c, 6),
               "usage": {k: usage.get(k) for k in ("prompt_tokens", "completion_tokens", "total_tokens", "finish_reason")}}
        if not text:
            rec["error"] = err
            return rec, None, tag, ""
        (raw_dir / f"{tag}.txt").write_text(text, encoding="utf-8")
        lab_list = [f"V{i+1}" for i in range(len(labels))]
        rank, perr, obj = parse_ranking(text, lab_list)
        if rank is None:
            rec["parse_error"] = perr
            return rec, None, tag, text
        # маппим метки → армы
        arm_rank = [labels[l] for l in rank]
        v = {"judge": jname, "chunk": f"{ch}/{ci}", "axis": axis, "rep": rep,
             "arm_ranking": arm_rank, "evidence_spans": obj.get("evidence_spans", [])}
        return rec, v, tag, text

    with ThreadPoolExecutor(max_workers=5) as ex:
        for rec, v, tag, _ in ex.map(one, tasks):
            jcosts.append(rec)
            if v:
                verdicts.append(v)
            spent_box[0] += rec["cost"]
    return verdicts, jcosts


def reparse():
    """Пере-парсит raw-файлы робастным парсером в verdicts.jsonl (без API-вызовов)."""
    keymap = json.loads((JOUT / "_JUDGE_KEY.json").read_text(encoding="utf-8"))
    pat = re.compile(r"^(J\d+_[a-z0-9]+)_(\d+)_(\d+)_(fidelity|style)_r(\d+)$")
    verdicts, ok, bad = [], 0, 0
    for f in sorted((JOUT / "raw").glob("*.txt")):
        m = pat.match(f.stem)
        if not m:
            continue
        jname, ch, ci, axis, rep = m.groups()
        labels = keymap[f"{ch}/{ci}"][axis][rep]
        lab_list = [f"V{i+1}" for i in range(len(labels))]
        rank, err, obj = parse_ranking(f.read_text(encoding="utf-8"), lab_list)
        if rank is None:
            bad += 1
            continue
        ok += 1
        verdicts.append({"judge": jname, "chunk": f"{ch}/{ci}", "axis": axis, "rep": int(rep),
                         "arm_ranking": [labels[l] for l in rank],
                         "evidence_spans": obj.get("evidence_spans", []) if isinstance(obj, dict) else []})
    with (JOUT / "verdicts.jsonl").open("w", encoding="utf-8") as f:
        for v in verdicts:
            f.write(json.dumps(v, ensure_ascii=False) + "\n")
    print(f"reparse: {ok} валидных, {bad} невалидных → verdicts.jsonl")


def main():
    if "--reparse" in sys.argv:
        reparse()
        return
    if "--repeats" in sys.argv:
        globals()["REPEATS"] = int(sys.argv[sys.argv.index("--repeats") + 1])
    chunks = json.loads((DIAG / "chunks.json").read_text(encoding="utf-8"))
    arms = arms_present()
    print(f"армы в панели: {arms}  (repeats={REPEATS})")
    JOUT.mkdir(parents=True, exist_ok=True)

    # ключ (label→arm) фиксируем ЗАРАНЕЕ для всех судей — свап на (чанк×ось×повтор)
    keymap = {}
    for ch, ci in TARGETS:
        keymap[f"{ch}/{ci}"] = {}
        for axis in AXES:
            keymap[f"{ch}/{ci}"][axis] = {}
            for rep in range(REPEATS):
                keymap[f"{ch}/{ci}"][axis][str(rep)] = build_labels(arms, ch, ci, axis, rep)
    (JOUT / "_JUDGE_KEY.json").write_text(json.dumps(keymap, ensure_ascii=False, indent=1), encoding="utf-8")

    all_v, all_c, spent_box = [], [], [0.0]
    cfp = ARMS / "judge_costs.jsonl"
    for jname, spec in JUDGES.items():
        if spent_box[0] > JUDGE_CAP:
            print(f"!! судейский кап ${JUDGE_CAP} достигнут (${spent_box[0]:.2f}) — {jname} пропущен"); continue
        t0 = time.time()
        v, c = run_judge(jname, spec, arms, chunks, keymap, spent_box)
        with cfp.open("a", encoding="utf-8") as f:
            for rec in c:
                f.write(json.dumps(rec, ensure_ascii=False) + "\n")
        ok = sum(1 for x in c if "error" not in x and "parse_error" not in x)
        print(f"{jname}: {ok}/{len(c)} валидных, {time.time()-t0:.0f}s, спенд-накопл ${spent_box[0]:.3f}")
        all_v += v; all_c += c

    with (JOUT / "verdicts.jsonl").open("w", encoding="utf-8") as f:
        for v in all_v:
            f.write(json.dumps(v, ensure_ascii=False) + "\n")
    print(f"\nвсего валидных вердиктов: {len(all_v)}  судейский спенд: ${spent_box[0]:.4f}")


if __name__ == "__main__":
    main()
