#!/usr/bin/env python3
"""Полигон, пакет-8: драйвер БОЕВОГО запроса роли терминолога.

Отправляет ровно то, что собрал `termharness -mode build` (системный промт роли + якорь ⟦TM-CANON⟧ +
блок кандидатов) — ни одного своего слова в промте нет. Полигон отвечает только за сетку параметров,
повторы и деньги.

Дисциплина (нормы D39.46/47):
  - СМЕТА печатается ДО первого вызова; `--dry-run` даёт её без единой траты;
  - сырьё каждого вызова сохраняется целиком (сообщения, ответ, usage, finish_reason) — выводы потом
    пересчитываются ИЗ сырья, а не из памяти прогона;
  - деньги считаются из фактического usage по `backend/configs/models.yaml` (read-only), с учётом
    кэш-цены prompt_tokens_details.cached_tokens;
  - max_tokens = 8000 — не круглое число, а пол `min_max_tokens` deepseek-v4-flash из models.yaml,
    к которому боевой `terminologyCallBudget` и приводит батч такого размера (est/2+256 < 8000);
  - thinking у DeepSeek НЕ отключается (гардрейл проекта: эхо-мина).
"""
from __future__ import annotations

import argparse
import json
import os
import sys
import time
from pathlib import Path

import yaml
from openai import OpenAI

REPO = Path(__file__).resolve().parents[2]


def load_prices(models_yaml: Path) -> dict:
    doc = yaml.safe_load(models_yaml.read_text(encoding="utf-8")) or {}
    out = {}
    for name, m in (doc.get("models") or {}).items():
        p = (m or {}).get("price") or {}
        out[name] = (p.get("input_per_m", 0.0), p.get("output_per_m", 0.0), p.get("cached_per_m", 0.0))
    return out


def main() -> int:
    ap = argparse.ArgumentParser()
    ap.add_argument("--req", required=True, type=Path, help="выход termharness -mode build")
    ap.add_argument("--raw-dir", required=True, type=Path)
    ap.add_argument("--tag", required=True, help="метка точки сетки (идёт в имена файлов сырья)")
    ap.add_argument("--repeats", type=int, default=1)
    ap.add_argument("--model", default="deepseek-v4-flash")
    ap.add_argument("--base-url", default="https://api.deepseek.com/v1")
    ap.add_argument("--key-env", default="DEEPSEEK_API_KEY")
    ap.add_argument("--max-tokens", type=int, default=8000)
    ap.add_argument("--dry-run", action="store_true")
    a = ap.parse_args()

    req = json.loads(a.req.read_text(encoding="utf-8"))
    batches = req["batches"]
    prices = load_prices(REPO / "backend/configs/models.yaml")
    pin, pout, pcached = prices.get(a.model, (0.0, 0.0, 0.0))

    # Оценка входа: тот же грубый делитель, что в пакете-7 (симв./3.2). Смета — верхняя граница по
    # выходу (max_tokens), как и боевая terminologyEstimateUSD, поэтому факт всегда ниже.
    chars = sum(len(m["content"]) for b in batches for m in b["messages"])
    est_in = chars / 3.2 * a.repeats
    est_out = len(batches) * a.max_tokens * a.repeats
    print(f"ПЛАН [{a.tag}]: батчей {len(batches)} × повторов {a.repeats} = {len(batches)*a.repeats} вызовов; "
          f"термов {req['candidates']}; kwic {req['kwic_per']}×{req['kwic_width']}")
    print(f"СМЕТА ДО ВЫЗОВОВ: вход ≈{est_in/1000:.1f}k ток → ${est_in/1e6*pin:.4f}; "
          f"выход ≤{est_out/1000:.1f}k ток → ≤${est_out/1e6*pout:.4f}; "
          f"ИТОГО ≤ ${est_in/1e6*pin + est_out/1e6*pout:.4f}")
    if a.dry_run:
        return 0

    key = os.environ.get(a.key_env)
    if not key:
        print(f"нет ключа {a.key_env}", file=sys.stderr)
        return 2
    client = OpenAI(api_key=key, base_url=a.base_url)
    a.raw_dir.mkdir(parents=True, exist_ok=True)

    total = 0.0
    for rep in range(a.repeats):
        for b in batches:
            # $0-резюм по образцу боевого чекпойнта: уже оплаченный батч не перепокупается. Без этого
            # любой таймаут харнесса означал бы повторную оплату всей точки сетки.
            done = a.raw_dir / f"{a.tag}_r{rep}_b{b['index']}.json"
            if done.exists() and "error" not in json.loads(done.read_text(encoding="utf-8")):
                continue
            msgs = [{"role": m["role"], "content": m["content"]} for m in b["messages"]]
            t0 = time.time()
            try:
                resp = client.chat.completions.create(model=a.model, messages=msgs, max_tokens=a.max_tokens)
            except Exception as e:  # noqa: BLE001 — ошибка вызова тоже улика, не повод падать молча
                print(f"[{a.tag} r{rep} b{b['index']}] ОШИБКА: {e}", file=sys.stderr)
                (a.raw_dir / f"{a.tag}_r{rep}_b{b['index']}.json").write_text(
                    json.dumps({"tag": a.tag, "rep": rep, "batch": b["index"], "error": str(e)},
                               ensure_ascii=False, indent=1), encoding="utf-8")
                continue
            dt = time.time() - t0
            txt = resp.choices[0].message.content or ""
            u = resp.usage
            cin = getattr(u, "prompt_tokens", 0)
            cout = getattr(u, "completion_tokens", 0)
            cached = getattr(getattr(u, "prompt_tokens_details", None), "cached_tokens", 0) or 0
            usd = (cin - cached) / 1e6 * pin + cached / 1e6 * pcached + cout / 1e6 * pout
            total += usd
            (a.raw_dir / f"{a.tag}_r{rep}_b{b['index']}.json").write_text(
                json.dumps({"tag": a.tag, "rep": rep, "batch": b["index"], "keys": b["keys"],
                            "messages": msgs, "response": txt,
                            "finish": resp.choices[0].finish_reason,
                            "usage": {"in": cin, "out": cout, "cached": cached},
                            "usd": usd, "seconds": round(dt, 1),
                            "kwic_per": req["kwic_per"], "kwic_width": req["kwic_width"]},
                           ensure_ascii=False, indent=1), encoding="utf-8")
            print(f"[{a.tag} r{rep} b{b['index']}] in={cin} (cached {cached}) out={cout} "
                  f"finish={resp.choices[0].finish_reason} ${usd:.5f} {dt:.0f}s")
    print(f"ИТОГО ФАКТ [{a.tag}]: ${total:.5f}")
    return 0


if __name__ == "__main__":
    sys.exit(main())