#!/usr/bin/env python3
"""exp15 — shared LLM call / cost / per-call-gate / ledger layer for paid rigs (judges, miner, arms).

Uses the openai SDK (proven in wirecheck.py) with the frozen live prices (freeze_facts.json) — NOT
exp14_common.PRICES (a stale 2026-07-10 mirror). Cost accounting is per-provider and VERIFIED against
the wire-check usage:
  grok-4.3 : total = prompt + completion + reasoning  => output billed = completion + reasoning (ADDITIVE)
  kimi-k2.6: verbose reasoning billed as output       => output billed = completion + reasoning (ADDITIVE)
  mistral  : no reasoning surcharge                    => output billed = completion
  glm-5    : output billed = completion; cached prompt tokens billed at cache_hit rate
  deepseek : output billed = completion (thinking folded / reported 0 for flash)
Tiered (grok/gemini) at 200k prompt tokens. Every paid call persists raw usage + cost (D30.10).
Per-call PREDICTED-cost gate BEFORE the call (worst-case max_tokens as output). Errors are returned,
not raised (honest accounting: record every attempt incl. failures — D37 lesson).
"""
from __future__ import annotations
import json
import os
import time
from pathlib import Path

from openai import OpenAI
from dotenv import load_dotenv

EVAL = Path(__file__).resolve().parent.parent
load_dotenv(EVAL / ".env")
FACTS = json.load(open("/home/ubuntu/books/gu-zhenren/exp15/freeze_facts.json"))
PRICES = FACTS["prices"]

# model -> provider wiring (base_url, env key, default temp, max_tokens floor, static extra_body,
#          whether reasoning tokens are billed additively as output)
MODELS = {
    "deepseek-v4-flash": dict(base="https://api.deepseek.com/v1", env="DEEPSEEK_API_KEY",
                              temp=0.3, max_tokens=8000, extra={}, reasoning_additive=False),
    "deepseek-v4-pro":   dict(base="https://api.deepseek.com/v1", env="DEEPSEEK_API_KEY",
                              temp=0.3, max_tokens=8000, extra={}, reasoning_additive=False),
    "glm-5":             dict(base="https://api.z.ai/api/paas/v4", env="ZAI_API_KEY",
                              temp=0.4, max_tokens=8000, extra={"thinking": {"type": "disabled"}},
                              reasoning_additive=False),
    "grok-4.3":          dict(base="https://api.x.ai/v1", env="XAI_API_KEY",
                              temp=0.0, max_tokens=8000, extra={}, reasoning_additive=True),  # reasoning ON (default)
    "mistral-large-latest": dict(base="https://api.mistral.ai/v1", env="MISTRAL_API_KEY",
                                 temp=0.0, max_tokens=8000, extra={"safe_prompt": False},
                                 reasoning_additive=False),
    "kimi-k2.6":         dict(base="https://api.moonshot.ai/v1", env="KIMI_API_KEY",
                              temp=1.0, max_tokens=16000, extra={}, reasoning_additive=True),  # temp MUST be 1
}


def _price(model, field, prompt_tokens=0):
    v = PRICES[model][field]
    if isinstance(v, list):
        tier_at = PRICES[model].get("tier_at", 200000)
        return v[1] if prompt_tokens >= tier_at else v[0]
    return v


def _reasoning_tokens(usage_obj) -> int:
    d = usage_obj if isinstance(usage_obj, dict) else getattr(usage_obj, "__dict__", {})
    if d.get("reasoning_tokens"):
        return d["reasoning_tokens"]
    ctd = d.get("completion_tokens_details")
    if ctd is not None:
        ctd = ctd if isinstance(ctd, dict) else getattr(ctd, "__dict__", {})
        return ctd.get("reasoning_tokens", 0) or 0
    # xAI: total = prompt + completion + reasoning -> infer if fields present
    return 0


def _cached_tokens(usage_obj) -> int:
    d = usage_obj if isinstance(usage_obj, dict) else getattr(usage_obj, "__dict__", {})
    # DeepSeek surfaces cache hits at top-level (prompt_cache_hit_tokens); GLM/OpenAI/Gemini nest it.
    if d.get("prompt_cache_hit_tokens"):
        return d["prompt_cache_hit_tokens"]
    ptd = d.get("prompt_tokens_details")
    if ptd is not None:
        ptd = ptd if isinstance(ptd, dict) else getattr(ptd, "__dict__", {})
        return ptd.get("cached_tokens", 0) or 0
    return 0


def cost(model, usage: dict) -> float:
    pt = usage.get("prompt_tokens", 0)
    ct = usage.get("completion_tokens", 0)
    rt = usage.get("reasoning_tokens", 0)
    cached = usage.get("cached_tokens", 0)
    pin = _price(model, "in", pt)
    pout = _price(model, "out", pt)
    phit = _price(model, "cache_hit", pt) if "cache_hit" in PRICES[model] else pin
    out_tok = ct + (rt if MODELS[model]["reasoning_additive"] else 0)
    in_cost = (pt - cached) * pin + cached * phit
    return (in_cost + out_tok * pout) / 1e6


# reasoning tokens are NOT capped by max_tokens for these models -> the gate must budget them too,
# else predict_cost under-estimates and a call can overshoot per_call_cap (kimi ~12k reasoning/para).
REASONING_CEIL = {"kimi-k2.6": 14000, "grok-4.3": 4000}


def predict_cost(model, in_tok, max_out) -> float:
    """Worst-case gate estimate: full max_tokens as billed output PLUS an uncapped-reasoning allowance
    for additive-reasoning models (kimi/grok), whose reasoning tokens are billed but not bounded by
    max_tokens. Without this the per-call/hard gate under-predicts (seen: kimi $0.124 vs $0.066)."""
    pin = _price(model, "in", in_tok)
    pout = _price(model, "out", in_tok)
    out_est = max_out + (REASONING_CEIL.get(model, 0) if MODELS[model]["reasoning_additive"] else 0)
    return (in_tok * pin + out_est * pout) / 1e6


class Spender:
    """Per-call predicted-cost gate + append-only jsonl ledger + running total (D30.10)."""

    def __init__(self, ledger_path: str, per_call_cap: float, hard_cap: float = None, seed_total: float = 0.0):
        # seed_total = spend already booked in OTHER ledgers, so hard_cap can be a GLOBAL ceiling even
        # though each block writes its own per-block ledger (review I8: per-block caps summed to 4x budget).
        self.path = Path(ledger_path)
        self.per_call_cap = per_call_cap
        self.hard_cap = hard_cap
        self.seed_total = seed_total
        self.total = seed_total
        if self.path.exists():
            for line in self.path.read_text(encoding="utf-8").splitlines():
                try:
                    self.total += json.loads(line).get("cost", 0.0)
                except Exception:
                    pass

    def gate(self, model, in_tok, max_out):
        pred = predict_cost(model, in_tok, max_out)
        if pred > self.per_call_cap:
            return False, pred, f"per-call gate: pred ${pred:.4f} > cap ${self.per_call_cap}"
        if self.hard_cap is not None and self.total + pred > self.hard_cap:
            return False, pred, f"hard-cap gate: total ${self.total:.4f}+${pred:.4f} > ${self.hard_cap}"
        return True, pred, "ok"

    def record(self, rec: dict):
        c = rec.get("cost", 0.0)
        self.total += c
        with self.path.open("a", encoding="utf-8") as f:
            f.write(json.dumps(rec, ensure_ascii=False) + "\n")


# proactive per-model min interval between calls (avoid 429 storms; mistral rate-limits aggressively).
# FLOOR full-chunk run measured mistral at 48% retry-fails @1.3s (max lat 64.7s) vs grok 0% @0.4s -> widen
# mistral spacing so its per-cell burst stays under the rate limit (fewer backoffs => cleaner AND faster).
MIN_INTERVAL = {"mistral-large-latest": 2.6, "grok-4.3": 0.4, "kimi-k2.6": 0.0}
_LAST_CALL = {}


def call(model, messages, temp=None, max_tokens=None, extra_body=None, timeout=240, retries=4):
    """Return {text, reasoning, usage, cost, model_returned, finish, err, latency_s}. Never raises.
    Retries transient errors (429 rate-limit, 5xx, timeouts) with exponential backoff (5/15/35/60s)."""
    m = MODELS[model]
    key = os.environ.get(m["env"])
    if not key:
        return {"err": f"no key {m['env']}", "cost": 0.0, "usage": {}}
    mi = MIN_INTERVAL.get(model, 0.0)
    if mi:
        dt = time.time() - _LAST_CALL.get(model, 0.0)
        if dt < mi:
            time.sleep(mi - dt)
    _LAST_CALL[model] = time.time()
    cli = OpenAI(base_url=m["base"], api_key=key, timeout=timeout)
    body_extra = dict(m["extra"])
    if extra_body:
        body_extra.update(extra_body)
    kw = dict(model=model, messages=messages,
              temperature=m["temp"] if temp is None else temp,
              max_tokens=max_tokens or m["max_tokens"])
    if body_extra:
        kw["extra_body"] = body_extra
    t0 = time.time()
    backoff = [5, 15, 35, 60]
    last_err = None
    retry_fails = 0     # transient attempts that were retried (may have billed server-side; review I8)
    for attempt in range(retries + 1):
        try:
            r = cli.chat.completions.create(**kw)
            break
        except Exception as e:  # noqa: BLE001
            last_err = f"{type(e).__name__}: {str(e)[:200]}"
            transient = any(s in last_err.lower() for s in ("ratelimit", "429", "rate limit", "timeout",
                                                            "timed out", "500", "502", "503", "529", "overloaded"))
            if attempt < retries and transient:
                retry_fails += 1
                time.sleep(backoff[min(attempt, len(backoff) - 1)])
                continue
            return {"err": last_err, "cost": 0.0, "usage": {}, "latency_s": round(time.time() - t0, 1),
                    "retry_fails": retry_fails}
    u = r.usage
    usage = {"prompt_tokens": getattr(u, "prompt_tokens", 0),
             "completion_tokens": getattr(u, "completion_tokens", 0),
             "total_tokens": getattr(u, "total_tokens", 0),
             "reasoning_tokens": _reasoning_tokens(u),
             "cached_tokens": _cached_tokens(u)}
    # xAI: reasoning often only inferable from total-prompt-completion
    if usage["reasoning_tokens"] == 0 and usage["total_tokens"] > usage["prompt_tokens"] + usage["completion_tokens"]:
        usage["reasoning_tokens"] = usage["total_tokens"] - usage["prompt_tokens"] - usage["completion_tokens"]
    msg = r.choices[0].message
    text = (msg.content or "").strip()
    reasoning = (getattr(msg, "reasoning_content", None) or "").strip()
    c = cost(model, usage)
    return {"text": text, "reasoning": reasoning, "usage": usage, "cost": round(c, 6),
            "model_returned": getattr(r, "model", None), "finish": r.choices[0].finish_reason,
            "err": None, "latency_s": round(time.time() - t0, 1), "retry_fails": retry_fails}


if __name__ == "__main__":
    # smoke: cost accounting sanity vs the frozen wire-check numbers
    grok_usage = {"prompt_tokens": 223, "completion_tokens": 11, "reasoning_tokens": 262, "total_tokens": 496, "cached_tokens": 0}
    print("grok additive cost (should exceed the wire-check's completion-only $0.000306):",
          round(cost("grok-4.3", grok_usage), 6))
    print("  = (223*1.25 + (11+262)*2.50)/1e6 =", round((223*1.25 + 273*2.50)/1e6, 6))
    mis_usage = {"prompt_tokens": 41, "completion_tokens": 15, "reasoning_tokens": 0, "total_tokens": 56, "cached_tokens": 0}
    print("mistral cost:", round(cost("mistral-large-latest", mis_usage), 6), "= (41*0.5+15*1.5)/1e6 =", round((41*0.5+15*1.5)/1e6, 6))
