#!/usr/bin/env python3
"""exp14 — общий харнес: провайдер-вызов, usage→$ (3 reasoning-режима), per-call
predicted-cost кап (НЕ group-level — урок exp13 +10%), append-only персист usage/cost,
токенайзер-счёт выходных токенов (кириллица-фертильность), spec-фабрика.

Зона eval/. Цены — ЗАМОРОЖЕННАЯ зеркальная копия backend/configs/models.yaml
(prices_checked 2026-07-10) + фронтир gpt-5.4 live-фактчек 2026-07-11
(developers.openai.com/api/docs/pricing). Единственный источник истины — models.yaml;
при расхождении верить Go/models.yaml. Ключи из eval/.env через load_env_file (НЕ читаю .env).
"""
from __future__ import annotations
import json, os, sys, time, urllib.request, urllib.error
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parent.parent))  # script now under eval/expNN/ — reach up to parent eval/ for the refusal_bench oracle
from refusal_bench import load_env_file  # noqa: E402

load_env_file()

DIAG = Path("/home/ubuntu/books/gu-zhenren/exp14")
DIAG.mkdir(parents=True, exist_ok=True)
RERUN = Path("/home/ubuntu/books/gu-zhenren/rerun")
TOKDIR = Path(__file__).resolve().parent / "tokenizers"

# ── Цены $/1M (in, out, cached_in) + reasoning-режим. Зеркало models.yaml. ──────────
PRICES = {
    "deepseek-v4-flash": (0.14, 0.28, 0.0028, "subset"),
    "deepseek-v4-pro": (0.435, 0.87, 0.003625, "subset"),
    "glm-5": (1.0, 3.2, 0.2, "subset"),
    "glm-5.1": (1.4, 4.4, 0.26, "subset"),
    "kimi-k2.6": (0.95, 4.0, 0.16, "subset"),
    "grok-4.3": (1.25, 2.50, 1.25, "additive"),
    "gemini-3.1-pro-preview": (2.0, 12.0, 0.20, "additive_total"),
    "gpt-5-mini": (0.25, 2.0, 0.025, "subset"),
    "gpt-5.4": (2.50, 15.0, 0.25, "subset"),  # live 2026-07-11
    "mistral-large-2512": (0.5, 1.5, 0.5, "subset"),
}
FALLBACK = "deepseek-v4-flash"  # никогда $0

# ── Спеки провайдеров (endpoint/keyenv/wire-квирки из 00-provider-quirks + models.yaml) ─
def spec(model: str, *, temp: float | None = None, max_tokens: int | None = None,
         reasoning_effort: str | None = None) -> dict:
    """Собирает per-call spec. temp/max_tokens/reasoning_effort перекрывают дефолты."""
    base_by = {
        "deepseek-v4-flash": ("https://api.deepseek.com/v1", "DEEPSEEK_API_KEY", {}),
        "deepseek-v4-pro": ("https://api.deepseek.com/v1", "DEEPSEEK_API_KEY", {}),
        "glm-5": ("https://api.z.ai/api/paas/v4", "ZAI_API_KEY", {"thinking": {"type": "disabled"}}),
        "glm-5.1": ("https://api.z.ai/api/paas/v4", "ZAI_API_KEY", {"thinking": {"type": "disabled"}}),
        "kimi-k2.6": ("https://api.moonshot.ai/v1", "KIMI_API_KEY", {}),
        "grok-4.3": ("https://api.x.ai/v1", "XAI_API_KEY", {}),
        "gemini-3.1-pro-preview": ("https://generativelanguage.googleapis.com/v1beta/openai", "GEMINI_API_KEY", {}),
        "gpt-5-mini": ("https://api.openai.com/v1", "OPENAI_API_KEY", {}),
        "gpt-5.4": ("https://api.openai.com/v1", "OPENAI_API_KEY", {}),
        "mistral-large-2512": ("https://api.mistral.ai/v1", "MISTRAL_API_KEY", {"safe_prompt": False}),
    }
    base, keyenv, extra = base_by[model]
    extra = dict(extra)
    s = {"model": model, "base": base, "keyenv": keyenv, "extra": extra}
    # reasoning wire per family
    if model in ("gpt-5-mini", "gpt-5.4"):
        s["openai_reasoner"] = True  # max_completion_tokens, no temperature
        extra["reasoning_effort"] = reasoning_effort or "medium"
    elif model == "kimi-k2.6":
        s["force_temp"] = 1.0
    elif model == "grok-4.3":
        if reasoning_effort:  # xAI plain reasoning_effort in body (none/low/medium/high)
            extra["reasoning_effort"] = reasoning_effort
    # temp
    if model in ("gpt-5-mini", "gpt-5.4"):
        s["temp"] = None
    elif model == "kimi-k2.6":
        s["temp"] = 1.0
    else:
        s["temp"] = 0.4 if temp is None else temp
    # max_tokens floors (D24.3)
    floor = 16000 if model == "kimi-k2.6" else (12000 if model == "gpt-5.4" else 8000)
    s["max_tokens"] = max_tokens or floor
    return s


def call(s: dict, messages: list[dict], timeout: int = 300) -> tuple[str, str | None, dict]:
    """Один chat/completions. Возврат (text, err|None, usage со стампом finish_reason).
    err — структурная строка: http|CODE|.. / transport|.. / content_filter|.. / empty|.."""
    key = os.environ.get(s["keyenv"], "")
    body = {"model": s["model"], "messages": messages, **s.get("extra", {})}
    if s.get("openai_reasoner"):
        body["max_completion_tokens"] = s["max_tokens"]  # gpt-5: max_tokens → 400
    else:
        body["max_tokens"] = s["max_tokens"]
    if s.get("temp") is not None:
        body["temperature"] = s["temp"]
    data = json.dumps(body).encode()
    req = urllib.request.Request(s["base"].rstrip("/") + "/chat/completions", data=data,
        headers={"Content-Type": "application/json", "Authorization": f"Bearer {key}"})
    try:
        with urllib.request.urlopen(req, timeout=timeout) as r:
            d = json.loads(r.read())
    except urllib.error.HTTPError as e:
        return "", f"http|{e.code}|{e.read()[:300].decode('utf-8','replace')}", {}
    except Exception as e:
        return "", f"transport|{type(e).__name__}|{str(e)[:200]}", {}
    ch = (d.get("choices") or [{}])[0]
    finish = ch.get("finish_reason", "")
    msg = ch.get("message", {}) or {}
    text = (msg.get("content") or "").strip()
    usage = d.get("usage", {}) or {}
    usage["finish_reason"] = finish
    if str(finish).lower() in ("content_filter", "prohibited_content", "safety") or "prohibited" in str(finish).lower():
        return "", f"content_filter|finish={finish}", usage
    if not text:
        has_reason = bool((msg.get("reasoning_content") or "").strip())
        return "", f"empty|finish={finish}|reasoning={has_reason}", usage
    return text, None, usage


def cost_of(model: str, usage: dict) -> float:
    """usage→$ с 3 reasoning-режимами (subset/additive/additive_total)."""
    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, _c, regime = PRICES.get(model, (*PRICES[FALLBACK][:3], "subset"))
    if regime == "additive":  # grok: reasoning_tokens отдельным полем, поверх completion
        reason = usage.get("reasoning_tokens", 0) or 0
        out = comp + reason
    elif regime == "additive_total":  # gemini: thinking только в total
        out = max(comp, total - inp) if total else comp
    else:  # subset: reasoning ⊆ completion
        out = comp
    return (inp * pin + out * pout) / 1_000_000


def predict_cost(model: str, in_tokens_est: int, max_out: int) -> float:
    """Верхняя оценка ДО вызова: in_est·pin + max_out·pout (потолок max_tokens как output)."""
    pin, pout, _c, _r = PRICES.get(model, (*PRICES[FALLBACK][:3], "subset"))
    return (in_tokens_est * pin + max_out * pout) / 1_000_000


# ── Токенайзеры (кириллица-фертильность; deepseek-v4/kimi отсутствуют → v3.2 прокси) ──
_TOK = {}
def _load_tok(name: str):
    from tokenizers import Tokenizer
    p = TOKDIR / name / "tokenizer.json"
    if name not in _TOK:
        _TOK[name] = Tokenizer.from_file(str(p))
    return _TOK[name]

def count_tokens(text: str, family: str = "glm") -> int:
    name = {"glm": "glm-4.6", "deepseek": "deepseek-v3.2"}.get(family, "glm-4.6")
    return len(_load_tok(name).encode(text, add_special_tokens=False).ids)


# ── Per-call кап + персист ────────────────────────────────────────────────────────
class Spender:
    """Per-key running spend + per-call predicted-cost кап (обязателен ДО вызова)."""
    def __init__(self, costs_path: Path, caps: dict[str, float]):
        self.path = costs_path
        self.caps = caps  # model → per-call cap $
        self.by_key: dict[str, float] = {}

    def guard(self, model: str, in_est: int, max_out: int) -> tuple[bool, float]:
        pc = predict_cost(model, in_est, max_out)
        cap = self.caps.get(model, 0.50)
        return (pc <= cap, pc)

    def record(self, rec: dict) -> float:
        c = cost_of(rec["model"], rec.get("usage", {}))
        rec["cost"] = round(c, 6)
        keyenv = rec.get("keyenv", rec["model"])
        self.by_key[keyenv] = self.by_key.get(keyenv, 0.0) + c
        with open(self.path, "a", encoding="utf-8") as f:
            f.write(json.dumps(rec, ensure_ascii=False) + "\n")
        return c


def messages_with_injection(system_prefix: str, injection: str, user: str) -> list[dict]:
    """Layout render.go MessagesWithInjection: system(prefix,cache) → system(inj) → user."""
    m = [{"role": "system", "content": system_prefix}]
    if injection.strip():
        m.append({"role": "system", "content": injection})
    m.append({"role": "user", "content": user})
    return m


if __name__ == "__main__":
    # смоук: цены/токенайзер/предикт
    ru = "Он шёл по тёмной улице, и снег ложился ему на плечи, укрывая следы."
    print("glm-4.6 tokens:", count_tokens(ru, "glm"))
    print("deepseek(v3.2) tokens:", count_tokens(ru, "deepseek"))
    u = {"prompt_tokens": 5000, "completion_tokens": 3000, "total_tokens": 8200}
    for m in ("glm-5", "gpt-5.4", "gemini-3.1-pro-preview", "grok-4.3"):
        print(f"cost_of {m}: ${cost_of(m, u):.5f}  predict(5k in, 8k out): ${predict_cost(m,5000,8000):.4f}")