#!/usr/bin/env python3
"""Локальный стенд TextMachine (задача 3 полигона, docs/research/06-local-models.md).

Меряет на установленных ollama-моделях:
  - скорость: prompt_eval tok/s и generation tok/s (из метрик /api/generate);
  - память: пик VRAM (nvidia-smi) и RSS ollama во время генерации;
  - качество: сохраняет переводы ja→ru / zh→ru для субъективного сравнения
    (эталон DeepSeek API добавляется тем же скриптом при наличии ключа).

Запуск: eval/.venv/bin/python eval/local_bench.py [--models m1,m2] [--num-ctx 8192]
Выход: eval/data/local_bench/<model>/<fragment>.txt + results.json + таблица в stdout.
"""
from __future__ import annotations

import argparse
import json
import os
import re
import subprocess
import threading
import time
import urllib.error
import urllib.request
from pathlib import Path

ROOT = Path(__file__).resolve().parent
OLLAMA = "http://localhost:11434"

# WSL: в env задан http_proxy, а NO_PROXY=<local> (WinINET-формат) urllib не понимает —
# localhost-запросы уходят на прокси и получают 403. Для ollama прокси отключаем явно.
_NO_PROXY_OPENER = urllib.request.build_opener(urllib.request.ProxyHandler({}))


def _load_env_file() -> None:
    """Ключи из eval/.env (gitignored, закрыт от Claude deny-правилом) или
    legacy ~/.textmachine.env; только в env процесса, не печатаются."""
    for cand in (Path(__file__).resolve().parent / ".env",
                 Path("~/.textmachine.env").expanduser()):
        if not cand.exists():
            continue
        for line in cand.read_text().splitlines():
            line = line.strip()
            if not line or line.startswith("#") or "=" not in line:
                continue
            k, _, v = line.partition("=")
            if v.strip():
                os.environ.setdefault(k.strip(), v.strip().strip('"').strip("'"))
        return


_load_env_file()

PROMPT_TMPL = (
    "Ты профессиональный литературный переводчик с {src_name} на русский. "
    "Переведи фрагмент художественной прозы. Требования: точность смысла, живой "
    "литературный русский без канцелярита и калек, сохранить все предложения и реплики. "
    "Выведи только перевод.\n\n{text}"
)
SRC_NAMES = {"ja": "японского", "zh": "китайского"}

# Фрагменты: (lang, файл, срез символов) — ~1.2-1.5k знаков, самодостаточные сцены
FRAGMENTS = [
    ("ja", "akutagawa-rashomon.txt", 0, 1400),
    ("ja", "dazai-hashire-merosu.txt", 0, 1400),
    ("zh", "luxun-zhufu.txt", 0, 1500),
]


def gpu_mem_mb() -> int:
    try:
        out = subprocess.run(
            ["nvidia-smi", "--query-gpu=memory.used", "--format=csv,noheader,nounits"],
            capture_output=True, text=True, timeout=5).stdout.strip()
        return int(out.splitlines()[0])
    except Exception:
        return -1


def ram_used_gb() -> float:
    with open("/proc/meminfo") as f:
        mi = {l.split(":")[0]: int(l.split()[1]) for l in f}
    return (mi["MemTotal"] - mi["MemAvailable"]) / 1e6


class PeakSampler(threading.Thread):
    def __init__(self):
        super().__init__(daemon=True)
        self.peak_vram = 0
        self.peak_ram = 0.0
        self.stop_flag = False

    def run(self):
        while not self.stop_flag:
            self.peak_vram = max(self.peak_vram, gpu_mem_mb())
            self.peak_ram = max(self.peak_ram, ram_used_gb())
            time.sleep(1)


def ollama_generate(model: str, prompt: str, num_ctx: int, timeout: int = 1800,
                    think: bool = False) -> dict:
    # think=False (дефолт): у моделей 3.5/GLM thinking включён по умолчанию и съедает
    # бюджет вывода на рассуждения (перевод обрезается/пустой) — для роли переводчика
    # отключаем. think=True: явно включаем + большой num_predict, чтобы модель успела
    # и порассуждать, и выдать перевод (замер влияния reasoning на качество).
    # num_predict — потолок вывода, чтобы модель не убегала.
    req = urllib.request.Request(
        f"{OLLAMA}/api/generate",
        data=json.dumps({
            "model": model, "prompt": prompt, "stream": False, "think": think,
            "options": {"num_ctx": num_ctx, "temperature": 0.3,
                        "num_predict": 6000 if think else 2048},
        }).encode(),
        headers={"Content-Type": "application/json"})
    try:
        with _NO_PROXY_OPENER.open(req, timeout=timeout) as resp:
            return json.loads(resp.read())
    except urllib.error.HTTPError as e:
        # не все модели принимают think=False (нет thinking-способности) — повтор без него
        if e.code == 400:
            body = json.loads(req.data)
            body.pop("think", None)
            req2 = urllib.request.Request(f"{OLLAMA}/api/generate", data=json.dumps(body).encode(),
                                          headers={"Content-Type": "application/json"})
            with _NO_PROXY_OPENER.open(req2, timeout=timeout) as resp:
                return json.loads(resp.read())
        raise


def strip_think(text: str) -> str:
    """Qwen3 в дефолте может отдавать <think>…</think> — вырезаем для оценки перевода."""
    return re.sub(r"<think>.*?</think>\s*", "", text, flags=re.S).strip()


def deepseek_generate(prompt: str) -> dict | None:
    key = os.environ.get("DEEPSEEK_API_KEY")
    if not key:
        return None
    req = urllib.request.Request(
        "https://api.deepseek.com/v1/chat/completions",
        data=json.dumps({"model": "deepseek-chat", "temperature": 0.3,
                         "messages": [{"role": "user", "content": prompt}]}).encode(),
        headers={"Content-Type": "application/json", "Authorization": f"Bearer {key}"})
    t0 = time.time()
    with urllib.request.urlopen(req, timeout=600) as resp:
        data = json.loads(resp.read())
    return {"response": data["choices"][0]["message"]["content"],
            "elapsed_s": time.time() - t0, "usage": data.get("usage", {})}


def installed_models() -> list[str]:
    with _NO_PROXY_OPENER.open(f"{OLLAMA}/api/tags", timeout=10) as resp:
        return [m["name"] for m in json.loads(resp.read())["models"]]


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("--models", help="через запятую; по умолчанию все установленные")
    ap.add_argument("--num-ctx", type=int, default=8192)
    ap.add_argument("--out", default=str(ROOT / "data" / "local_bench"))
    ap.add_argument("--think", action="store_true",
                    help="включить thinking-режим (reasoning) + большой num_predict — замер влияния на качество")
    args = ap.parse_args()

    models = args.models.split(",") if args.models else installed_models()
    out_dir = Path(args.out)
    out_dir.mkdir(parents=True, exist_ok=True)

    results = []
    for model in models:
        print(f"\n=== {model} (num_ctx={args.num_ctx}) ===")
        subprocess.run(["ollama", "stop", model.split(":")[0]], capture_output=True)
        base_vram, base_ram = gpu_mem_mb(), ram_used_gb()
        for lang, fname, a, b in FRAGMENTS:
            src_path = ROOT / "data" / "samples" / lang / fname
            if not src_path.exists():
                print(f"  [skip] нет {src_path}")
                continue
            text = src_path.read_text()[a:b]
            prompt = PROMPT_TMPL.format(src_name=SRC_NAMES[lang], text=text)
            sampler = PeakSampler()
            sampler.start()
            t0 = time.time()
            try:
                r = ollama_generate(model, prompt, args.num_ctx, think=args.think)
            except Exception as e:
                sampler.stop_flag = True
                print(f"  {fname}: ОШИБКА {type(e).__name__}: {str(e)[:200]}")
                results.append({"model": model, "fragment": fname, "error": str(e)[:500]})
                continue
            elapsed = time.time() - t0
            sampler.stop_flag = True
            sampler.join(2)
            gen_tps = r.get("eval_count", 0) / max(r.get("eval_duration", 1), 1) * 1e9
            pp_tps = r.get("prompt_eval_count", 0) / max(r.get("prompt_eval_duration", 1), 1) * 1e9
            translation = strip_think(r.get("response", ""))
            mdir = out_dir / model.replace("/", "_").replace(":", "_")
            mdir.mkdir(exist_ok=True)
            (mdir / f"{fname.removesuffix('.txt')}.ru.txt").write_text(translation)
            rec = {"model": model, "fragment": fname, "lang": lang,
                   "src_chars": len(text),
                   "prompt_eval_tps": round(pp_tps, 1), "gen_tps": round(gen_tps, 1),
                   "gen_tokens": r.get("eval_count"), "elapsed_s": round(elapsed, 1),
                   "peak_vram_mb": sampler.peak_vram, "base_vram_mb": base_vram,
                   "peak_ram_gb": round(sampler.peak_ram, 1),
                   "base_ram_gb": round(base_ram, 1),
                   "out_chars": len(translation)}
            results.append(rec)
            print(f"  {fname}: gen {gen_tps:.1f} tok/s, prompt {pp_tps:.0f} tok/s, "
                  f"{elapsed:.0f}s, VRAM пик {sampler.peak_vram} MB, RAM пик {sampler.peak_ram:.1f} GB")

    # эталон DeepSeek (если есть ключ)
    if os.environ.get("DEEPSEEK_API_KEY"):
        print("\n=== deepseek-chat (эталон API) ===")
        mdir = out_dir / "deepseek-api"
        mdir.mkdir(exist_ok=True)
        for lang, fname, a, b in FRAGMENTS:
            src_path = ROOT / "data" / "samples" / lang / fname
            text = src_path.read_text()[a:b]
            r = deepseek_generate(PROMPT_TMPL.format(src_name=SRC_NAMES[lang], text=text))
            (mdir / f"{fname.removesuffix('.txt')}.ru.txt").write_text(r["response"])
            results.append({"model": "deepseek-api", "fragment": fname,
                            "elapsed_s": round(r["elapsed_s"], 1), "usage": r["usage"]})
            print(f"  {fname}: {r['elapsed_s']:.0f}s")
    else:
        print("\n[note] DEEPSEEK_API_KEY нет — эталонные переводы не сняты")

    (out_dir / "results.json").write_text(json.dumps(results, ensure_ascii=False, indent=1))
    print(f"\nJSON: {out_dir}/results.json; переводы: {out_dir}/<model>/*.ru.txt")


if __name__ == "__main__":
    main()