#!/usr/bin/env python3
"""Экстракция терминов — БЕНЧМАРК локальных vs облачных моделей (write-path банка памяти, реестр G1).

Локалка касается банка памяти в 2 точках: эмбеддинги (retrieval_bench.py — bge-m3) и
ЭКСТРАКЦИЯ кандидатов в глоссарий из сырого текста (bootstrap write-path). Полигон exp03
мерил локалки как ПЕРЕВОДЧИКОВ (не годятся); экстракция — другая, более лёгкая задача.
Здесь мерим её нормально: много моделей, много кейсов, разные языки/эдж-кейсы, и КАЖДЫЙ
ответ верифицируется на здоровость (пустой/echo/config-error → ретрай и явная пометка
FAILED, НЕ засчитывается как recall 0 — урок предыдущего обрывочного прогона).

Запросы к провайдерам — через `refusal_bench.call_provider` + канонический `providers.json`
(квирки уже верны: deepseek thinking-ON+max_tokens 8000, kimi temp=1/24k, gpt-5-mini
max_completion_tokens+no-temp, gemini/glm thinking-off). Локаль — ollama native (think:false).
Плюс вариант deepseek-nothink (`thinking:disabled`) — дешёвая экстракция без reasoning.

Метрики (ДЕТЕРМИНИРОВАННЫЕ, без судьи — чтобы «верифицированно и точно»):
  recall     — доля эталонных сущностей, найденных (спот)
  render_ok  — доля найденных gold, где предложенный RU совпал с каноном (Рогов/Поливанов/Палладий/устоявш.)
  halluc     — доля извлечённого, чего НЕТ в исходном тексте (явный FP, детерминированно)
  health     — ok / empty / http / refused / unparseable (ретраи; в агрегат идут только ok)

Запуск: eval/.venv/bin/python eval/extract_bench.py [--models a,b] [--cases zh,ja,en]
Выход: eval/data/extract_bench/results.json + таблица.
"""
from __future__ import annotations

import argparse
import importlib.util
import json
import os
import re
import time
import unicodedata
import urllib.request
from pathlib import Path

from dotenv import dotenv_values

ROOT = Path(__file__).resolve().parent
OUT = ROOT / "data" / "extract_bench"
ENV = dotenv_values(ROOT / ".env")
for k, v in ENV.items():          # call_provider читает ключи из os.environ
    os.environ.setdefault(k, v or "")
NOPROXY = urllib.request.build_opener(urllib.request.ProxyHandler({}))

# reuse verified request machinery
_spec = importlib.util.spec_from_file_location("rb", ROOT / "refusal_bench.py")
rb = importlib.util.module_from_spec(_spec); _spec.loader.exec_module(rb)
PROVIDERS = {p["name"]: p for p in json.loads((ROOT / "providers.json").read_text())["providers"]}

CLOUD = ["deepseek", "grok", "glm", "kimi", "gemini", "openai"]
CLOUD_EXTRA = {  # deepseek без thinking — дешёвая экстракция (echo бьёт перевод, не JSON?)
    "deepseek-nothink": dict(PROVIDERS["deepseek"], name="deepseek-nothink",
                             extra_body={"thinking": {"type": "disabled"}}),
}
LOCAL = ["huihui_ai/qwen3-abliterated:8b", "huihui_ai/qwen3.5-abliterated:9b",
         "qwen3:8b", "qwen3-vojo:latest", "ruadapt-qwen3:8b",
         "huihui_ai/qwen3-abliterated:30b-a3b"]

SYS = "Ты — терминолог. Извлекаешь из текста имена собственные и культурные реалии для глоссария художественного перевода на русский."
USER = ("Извлеки из фрагмента ВСЕ имена собственные (персонажи, места, организации) и "
        "культурно-специфичные реалии/титулы. Для каждого верни: src (ровно как в тексте), "
        "type (character|place|org|realia|title), dst (перевод/транскрипция на русский по канону). "
        'Верни СТРОГО JSON без пояснений: {"terms":[{"src":"...","type":"...","dst":"..."}]}.\n\nФрагмент:\n<<<TEXT>>>')

# ---- корпус: реальные PD-чанки, разные языки/эдж-кейсы, gold с каноническим RU (accept-регэксп) ----
CASES = [
    # zh — имена + реалии + алиас
    dict(id="zh-ahq", lang="zh",
         text="我要给阿Q做正传，已经不止一两年了。那是赵太爷的儿子进了秀才的时候，锣声铛铛的报到村里来，阿Q正喝了两碗黄酒，因为他和赵太爷原来是本家。",
         gold={"阿Q": r"А[-\s]?[Кк]ью", "赵太爷": r"Чжао", "秀才": r"сюцай"}),
    dict(id="zh-zhufu", lang="zh",
         text="接着一声钝响，是送灶的爆竹。我是正在这一夜回到我的故乡鲁镇的，暂寓在鲁四老爷的宅子里，应该称之曰「四叔」。家中都在准备着「祝福」，这是鲁镇年终的大典，迎接福神。",
         gold={"送灶": r"очаг|Цзао", "鲁镇": r"Лучжэн", "鲁四老爷": r"дядюшк|[Лл]у\b|Лу[ -]", "祝福": r"жертвоприношен|благословени|[Мм]олени|счасть"}),
    dict(id="zh-xianglin", lang="zh",
         text="况且，一想到昨天遇见祥林嫂的事，也就使我不能安住。我在鲁镇所见的人们中，改变之大，可以说无过于她的了；她分明已经纯乎是一个乞丐了。",
         gold={"祥林嫂": r"Сянлин", "鲁镇": r"Лучжэн"}),
    dict(id="zh-wanghu", lang="zh",  # алиас: 王胡 = 王癞胡
         text="他看见王胡在那里赤着膊捉虱子。这王胡，又癞又胡，别人都叫他王癞胡，阿Q却删去了一个癞字，然而非常渺视他。",
         gold={"王胡": r"Ван|Бородат", "阿Q": r"А[-\s]?[Кк]ью"}),
    # ja — имена + реалии
    dict(id="ja-rashomon", lang="ja",
         text="一人の下人が、羅生門の下で雨やみを待っていた。羅生門が、朱雀大路にある以上は、この男のほかにも、雨やみをする者がありそうなものである。",
         gold={"下人": r"слуг|гэнин", "羅生門": r"Рас[её]мон", "朱雀大路": r"Судзаку"}),
    dict(id="ja-kyoto", lang="ja",
         text="この二三年、京都には、地震とか辻風とか火事とか饑饉とか云う災がつづいて起った。そこで洛中のさびれ方は一通りではない。",
         gold={"京都": r"Киото", "洛中": r"Ракутю|Лочжун|столиц|город"}),
    # ja — западные имена через катакану (translit_policy эдж-кейс)
    dict(id="ja-melos", lang="ja",
         text="メロスは激怒した。メロスは、村の牧人である。きょう未明メロスは村を出発し、十里はなれた此のシラクスの市にやって来た。",
         gold={"メロス": r"Мелос", "シラクス": r"Сиракуз"}),
    dict(id="ja-seri", lang="ja",  # длинное западное имя через катакану
         text="メロスには竹馬の友があった。セリヌンティウスである。今は此のシラクスの市で、石工をしている。",
         gold={"セリヌンティウス": r"Селинунти", "メロス": r"Мелос", "シラクス": r"Сиракуз"}),
    # ja — буддийское имя + топоним (эталон Стругацкого)
    dict(id="ja-hana", lang="ja",
         text="禅智内供の鼻と云えば、池尾で知らない者はない。長さは五六寸あって上唇の上から顋の下まで下っている。",
         gold={"禅智内供": r"Дзэн[тч]и|Найгу", "池尾": r"Икэ"}),
    # en — западные имена → русская транскрипция (render-эдж)
    dict(id="en-carter", lang="en",
         text="My name is John Carter; I am better known as Captain Jack Carter of Virginia. I spent nearly a year prospecting in company with another Confederate officer, Captain James K. Powell of Richmond, in an Arizona cave.",
         gold={"John Carter": r"Джон\s?Картер", "Virginia": r"Виргини|Вирджини", "Powell": r"Пауэл", "Richmond": r"Ричмонд", "Arizona": r"Аризон"}),
    dict(id="en-conan", lang="en",
         text="'I am Conan, a Cimmerian,' he answered. 'I came into Argos seeking employment, but with no wars forward, there was nothing to which I might turn my hand.'",
         gold={"Conan": r"Конан", "Cimmerian": r"[Кк]иммери", "Argos": r"Аргос"}),
    dict(id="en-wells", lang="en",
         text="The Time Traveller (for so it will be convenient to speak of him) was expounding a recondite matter to us. 'You must follow me carefully,' said Filby.",
         gold={"Time Traveller": r"[Пп]утешественник", "Filby": r"Филби"}),
]


def norm(s):
    return unicodedata.normalize("NFKC", str(s)).lower()


def parse_terms(txt):
    txt = re.sub(r"<think>.*?</think>\s*", "", txt or "", flags=re.S)
    m = re.search(r"\{.*\}|\[.*\]", txt, re.S)
    if not m:
        return None                      # unparseable
    try:
        d = json.loads(m.group(0))
    except Exception:
        return None
    terms = d.get("terms", d) if isinstance(d, dict) else d
    if not isinstance(terms, list):
        return []
    return [{"src": str(t["src"]), "type": str(t.get("type", "")), "dst": str(t.get("dst", ""))}
            for t in terms if isinstance(t, dict) and t.get("src")]


def call_local(model, text, max_pred=4000):
    payload = dict(model=model, stream=False, think=False, format="json",
                   options=dict(num_ctx=8192, temperature=0.2, num_predict=max_pred),
                   messages=[{"role": "system", "content": SYS},
                             {"role": "user", "content": USER.replace("<<<TEXT>>>", text)}])
    req = urllib.request.Request("http://localhost:11434/api/chat",
                                 data=json.dumps(payload).encode(), headers={"Content-Type": "application/json"})
    with NOPROXY.open(req, timeout=600) as r:
        d = json.load(r)
    return (d.get("message", {}) or {}).get("content", "") or "", d.get("done_reason", "")


def healthy_extract(kind, model, case):
    """Возвращает (terms|None, health, dt). Ретраит пустой ответ с бо́льшим бюджетом."""
    text = case["text"]
    for attempt in range(3):
        t0 = time.time()
        try:
            if kind == "local":
                raw, fin = call_local(model, text, max_pred=4000 * (attempt + 1))
                err = None if raw.strip() else f"empty|finish={fin}"
            else:
                pmap = CLOUD_EXTRA.get(model) or PROVIDERS[model]
                if attempt:  # ретрай пустого — поднять бюджет
                    pmap = dict(pmap, max_tokens=int(pmap.get("max_tokens", 8000) * 1.7))
                raw, err, _ = rb.call_provider(pmap, SYS, USER.replace("<<<TEXT>>>", text))
                raw = raw or ""
        except Exception as e:
            return None, f"transport:{type(e).__name__}", round(time.time() - t0, 1)
        dt = round(time.time() - t0, 1)
        if err:
            kindn = err.split("|", 1)[0]
            if kindn == "content_filter":
                return None, "refused", dt
            if kindn in ("empty",):
                continue                 # ретрай
            return None, f"http:{err[:40]}", dt
        terms = parse_terms(raw)
        if terms is None:
            if attempt < 2:
                continue                 # unparseable → ретрай
            return None, "unparseable", dt
        return terms, "ok", dt
    return None, "empty", dt


def score(case, terms):
    ntext = norm(case["text"])
    found, render = {}, {}
    for g, accept in case["gold"].items():
        ng = norm(g)
        hit = next((t for t in terms if ng in norm(t["src"]) or (norm(t["src"]) in ng and len(norm(t["src"])) >= 2)), None)
        found[g] = hit is not None
        render[g] = bool(hit and re.search(accept, hit["dst"]))
    ex = [t["src"] for t in terms]
    halluc = [s for s in ex if norm(s) and norm(s) not in ntext]
    return dict(recall=sum(found.values()) / len(found),
                render_ok=(sum(render.values()) / max(1, sum(found.values())) if any(found.values()) else 0.0),
                halluc_rate=round(len(halluc) / len(ex), 2) if ex else 0.0,
                n=len(terms), miss=[g for g, ok in found.items() if not ok],
                bad_render=[g for g in found if found[g] and not render[g]],
                halluc=halluc[:5], raw=[f"{t['src']}→{t['dst']}" for t in terms][:14])


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--models", default="")
    ap.add_argument("--cases", default="")
    args = ap.parse_args()
    OUT.mkdir(parents=True, exist_ok=True)
    cases = [c for c in CASES if not args.cases or c["lang"] in args.cases.split(",")]
    models = ([("local", m) for m in LOCAL] + [("cloud", m) for m in CLOUD]
              + [("cloud", "deepseek-nothink")])
    if args.models:
        want = set(args.models.split(","))
        models = [(k, m) for k, m in models if m in want]

    jsonl = OUT / "records.jsonl"
    done = set()
    if jsonl.exists():
        for ln in jsonl.read_text().splitlines():
            if ln.strip():
                r = json.loads(ln); done.add((r["model"], r["case"]))
    fh = jsonl.open("a", encoding="utf-8")
    records = [json.loads(ln) for ln in (jsonl.read_text().splitlines() if jsonl.exists() else []) if ln.strip()]

    for kind, model in models:
        print(f"\n### {model} ({kind})")
        for case in cases:
            if (model, case["id"]) in done:
                continue
            terms, health, dt = healthy_extract(kind, model, case)
            rec = dict(model=model, kind=kind, case=case["id"], lang=case["lang"], health=health, dt=dt)
            if health == "ok":
                rec.update(score(case, terms))
            records.append(rec)
            fh.write(json.dumps(rec, ensure_ascii=False) + "\n"); fh.flush()
            if health == "ok":
                print(f"  {case['id']:<12} recall={rec['recall']:.2f} render={rec['render_ok']:.2f} "
                      f"halluc={rec['halluc_rate']} n={rec['n']} {dt}s"
                      + (f" miss={rec['miss']}" if rec['miss'] else "")
                      + (f" bad_render={rec['bad_render']}" if rec['bad_render'] else ""))
            else:
                print(f"  {case['id']:<12} HEALTH={health} {dt}s")
    fh.close()

    # ---- агрегаты: только health==ok идут в recall/render/halluc; health-счётчики отдельно ----
    summary = {}
    for kind, model in models:
        recs = [r for r in records if r["model"] == model]
        ok = [r for r in recs if r["health"] == "ok"]
        health_ct = {}
        for r in recs:
            health_ct[r["health"]] = health_ct.get(r["health"], 0) + 1
        if ok:
            summary[model] = dict(
                kind=kind, n_ok=len(ok), n_total=len(recs), health=health_ct,
                recall=round(sum(r["recall"] for r in ok) / len(ok), 3),
                render_ok=round(sum(r["render_ok"] for r in ok) / len(ok), 3),
                halluc=round(sum(r["halluc_rate"] for r in ok) / len(ok), 3),
                avg_n=round(sum(r["n"] for r in ok) / len(ok), 1),
                avg_s=round(sum(r["dt"] for r in ok) / len(ok), 1))
        else:
            summary[model] = dict(kind=kind, n_ok=0, n_total=len(recs), health=health_ct)

    print("\n=== СВОДКА (экстракция; агрегат по health==ok) ===")
    print(f"  {'model':<34}{'kind':<7}{'recall':>7}{'render':>7}{'halluc':>7}{'n_ex':>5}{'sec':>6}{'ok/all':>8}")
    for m, s in summary.items():
        ratio = f"{s['n_ok']}/{s['n_total']}"
        if s["n_ok"]:
            print(f"  {m:<34}{s['kind']:<7}{s['recall']:>7}{s['render_ok']:>7}{s['halluc']:>7}"
                  f"{s['avg_n']:>5}{s['avg_s']:>6}{ratio:>8}")
        else:
            print(f"  {m:<34}{s['kind']:<7}{'—':>7}{'—':>7}{'—':>7}{'—':>5}{'—':>6}{ratio:>8}  health={s['health']}")
    print("  (health!=ok — пустой/echo/http/unparseable — НЕ засчитан как recall 0)")
    (OUT / "results.json").write_text(json.dumps(dict(summary=summary, records=records), ensure_ascii=False, indent=2), encoding="utf-8")
    print(f"\nsaved → {OUT/'results.json'}")


if __name__ == "__main__":
    main()