#!/usr/bin/env python3
"""
bench-decision-layer.py — measure cost and latency for one bounded decision,
four ways.

Backends
  jev        TypeSafe Jev decisions API            (needs TYPESAFE_API_KEY)
  json       LLM structured output / JSON mode     (needs OPENAI_API_KEY or BASE_URL)
  logit      Single constrained token via logit_bias on the same LLM
  encoder    Local fine-tuned encoder classifier   (needs torch + transformers)

Every backend answers the SAME labelled cases and is scored against the SAME
ground truth, so accuracy, p50/p95 latency and cost-per-decision are directly
comparable. Ground truth comes from the labels in the input file — not from
another model's opinion. That distinction is the whole point of this script.

Usage
  python scripts/bench-decision-layer.py \
      --cases data/decision-cases.jsonl \
      --backends json,logit,encoder \
      --out public/data/decision-layer-bench.csv

Input format (JSONL, one case per line)
  {"id": "t-001", "state": "Card declined three times...", "label": "billing"}

Prices are read from PRICES below and are point-in-time. Re-verify before
publishing any figure derived from this script.
"""

from __future__ import annotations

import argparse
import json
import os
import statistics
import sys
import time
from dataclasses import dataclass, field
from typing import Callable

import httpx


def _load_dotenv() -> None:
    """Pull keys from .env.local / .env into os.environ without echoing them."""
    for f in (".env.local", ".env"):
        if not os.path.exists(f):
            continue
        for line in open(f):
            line = line.strip()
            if not line or line.startswith("#") or "=" not in line:
                continue
            k, v = line.removeprefix("export ").split("=", 1)
            os.environ.setdefault(k.strip(), v.strip().strip("\"'"))


_load_dotenv()

# --- Label space under test -------------------------------------------------
# Default: 4-way support routing. Override with --labels <file.json> to run a
# public human-labelled set (e.g. AG News) — see data/decision-labels.json.
TASK = "Route this support ticket to one team."
LABELS: dict[str, str] = {
    "billing": "Payments, invoicing, refunds, card failures, subscription changes",
    "technical": "Bugs, outages, errors, integrations, API failures, data problems",
    "sales": "Pricing questions, upgrades, new accounts, demos, contract terms",
    "other": "Anything that fits none of the above",
}

# --- Prices, USD per 1M tokens -------------------------------------------------
# OpenAI: Standard tier, verified 2026-09-21 against developers.openai.com/api/docs/pricing
# (primary source — a secondary aggregator had gpt-5.6-sol wrong at $5/$30).
# Jev: TypeSafe's published early-access rate; output free; expected to change.
MODEL_PRICES: dict[str, dict[str, float]] = {
    "gpt-4.1-nano": {"input": 0.10, "output": 0.40},
    "gpt-4.1-mini": {"input": 0.40, "output": 1.60},
    "gpt-4o-mini": {"input": 0.15, "output": 0.60},
    "gpt-5.4-nano": {"input": 0.20, "output": 1.25},
    "gpt-5.6-luna": {"input": 0.20, "output": 1.20},
    "gpt-5.6-terra": {"input": 2.00, "output": 12.00},
    "gpt-5.6-sol": {"input": 4.00, "output": 20.00},
}
PRICES: dict[str, dict[str, float]] = {
    "jev": {"input": 0.042, "output": 0.0},
    "encoder": {"input": 0.0, "output": 0.0},  # amortised separately, see --gpu-hourly
}


def is_gpt5(model: str) -> bool:
    """GPT-5.x / o-series reject max_tokens, temperature and logit_bias."""
    return model.startswith(("gpt-5", "o1", "o3", "o4"))


@dataclass
class Result:
    backend: str
    case_id: str
    predicted: str | None
    expected: str
    latency_ms: float
    input_tokens: int = 0
    output_tokens: int = 0
    error: str | None = None
    confidence: float | None = None

    @property
    def correct(self) -> bool:
        return self.predicted == self.expected


@dataclass
class Backend:
    name: str
    run: Callable[[str], tuple[str | None, int, int, float | None]]
    notes: str = ""
    warmup: int = 1
    extra: dict = field(default_factory=dict)


# ---------------------------------------------------------------------------
# 1. Jev — one call, typed answer, probabilities over the label set.
# ---------------------------------------------------------------------------
def make_jev(client: httpx.Client, model: str = "jev-latest") -> Backend:
    key = os.environ.get("TYPESAFE_API_KEY")
    if not key:
        raise SystemExit("jev backend needs TYPESAFE_API_KEY")
    url = os.environ.get("TYPESAFE_BASE_URL", "https://api.typesafe.ai/v1/decisions")

    def run(state: str):
        r = client.post(
            url,
            headers={"Authorization": f"Bearer {key}"},
            json={
                "model": model,
                "state": state,
                "questions": {
                    "route": {
                        "type": "choice",
                        "instructions": TASK,
                        "criteria": LABELS,
                    }
                },
            },
            timeout=30.0,
        )
        r.raise_for_status()
        body = r.json()
        ans = body["answers"]["route"]
        usage = body.get("usage", {})
        return (
            ans["choice"],
            usage.get("input_tokens", 0),
            usage.get("output_tokens", 0),
            ans.get("confidence"),
        )

    return Backend("jev", run, notes="typed choice, probabilities returned")


# ---------------------------------------------------------------------------
# 2. JSON mode — the status quo. Sequential decode of a JSON object, then a
#    parse step that can fail, then a validation step that can also fail.
# ---------------------------------------------------------------------------
def make_json(client: httpx.Client, model: str) -> Backend:
    key = os.environ.get("OPENAI_API_KEY", "not-needed-for-local")
    base = os.environ.get("OPENAI_BASE_URL", "https://api.openai.com/v1")

    schema = {
        "type": "object",
        "properties": {"route": {"type": "string", "enum": list(LABELS)}},
        "required": ["route"],
        "additionalProperties": False,
    }
    prompt = (
        TASK + "\n\n"
        + "\n".join(f"- {k}: {v}" for k, v in LABELS.items())
        + "\n\nRespond with JSON matching the schema."
    )

    def run(state: str):
        r = client.post(
            f"{base}/chat/completions",
            headers={"Authorization": f"Bearer {key}"},
            json={
                "model": model,
                "messages": [
                    {"role": "system", "content": prompt},
                    {"role": "user", "content": state},
                ],
                "response_format": {
                    "type": "json_schema",
                    "json_schema": {"name": "route", "schema": schema, "strict": True},
                },
                **({"max_completion_tokens": 256} if is_gpt5(model)
                   else {"temperature": 0, "max_tokens": 32}),
            },
            timeout=120.0,
        )
        r.raise_for_status()
        body = r.json()
        usage = body.get("usage", {})
        raw = body["choices"][0]["message"]["content"]
        try:
            parsed = json.loads(raw)
            choice = parsed.get("route")
        except json.JSONDecodeError:
            choice = None  # the failure mode Jev's type safety removes
        if choice not in LABELS:
            choice = None
        return (
            choice,
            usage.get("prompt_tokens", 0),
            usage.get("completion_tokens", 0),
            None,
        )

    return Backend("json", run, notes="json_schema strict, parse can still fail")


# ---------------------------------------------------------------------------
# 3. logit_bias — the "no moat" baseline. Force exactly one token from a
#    hand-built alphabet, read the logprob as a calibrated-ish confidence.
#    This is the cheapest thing that answers the same question, and it is the
#    comparison TypeSafe's own benchmark does not run.
# ---------------------------------------------------------------------------
def make_logit(client: httpx.Client, model: str) -> Backend:
    key = os.environ.get("OPENAI_API_KEY", "not-needed-for-local")
    base = os.environ.get("OPENAI_BASE_URL", "https://api.openai.com/v1")

    if is_gpt5(model):
        # Measured 2026-09-21: every GPT-5.x model tested returns 400
        # "Unsupported parameter: 'logit_bias'". The trick needs an older
        # hosted model or a self-hosted one (vLLM / llama.cpp guided choice).
        raise SystemExit(f"logit backend: {model} does not support logit_bias")
    try:
        import tiktoken
    except ImportError:
        raise SystemExit("logit backend needs tiktoken: pip install tiktoken")

    # Map each label to a single-token letter. Single-token is the requirement;
    # the letters are arbitrary and must be resolved against the real encoding
    # rather than assumed — token ids differ per model family.
    letters = {lab: chr(ord("A") + i) for i, lab in enumerate(LABELS)}
    by_letter = {v: k for k, v in letters.items()}
    try:
        enc = tiktoken.encoding_for_model(model)
    except KeyError:
        enc = tiktoken.get_encoding("o200k_base")

    bias = {}
    for letter in by_letter:
        ids = enc.encode(letter)
        if len(ids) != 1:
            raise SystemExit(f"letter {letter!r} is not single-token under this encoding")
        bias[str(ids[0])] = 100

    prompt = (
        TASK + " Answer with exactly one letter.\n\n"
        + "\n".join(f"{letters[k]} = {k}: {v}" for k, v in LABELS.items())
    )

    def run(state: str):
        r = client.post(
            f"{base}/chat/completions",
            headers={"Authorization": f"Bearer {key}"},
            json={
                "model": model,
                "messages": [
                    {"role": "system", "content": prompt},
                    {"role": "user", "content": state},
                    # Prefill removes the model's urge to preamble. One token out.
                    {"role": "assistant", "content": "Answer:"},
                ],
                "logit_bias": bias,
                "max_tokens": 1,
                "temperature": 0,
                "logprobs": True,
                "top_logprobs": max(8, 2 * len(LABELS)),
            },
            timeout=60.0,
        )
        r.raise_for_status()
        body = r.json()
        usage = body.get("usage", {})
        choice_obj = body["choices"][0]
        letter = (choice_obj["message"]["content"] or "").strip()[:1]

        conf = None
        try:
            import math

            top = choice_obj["logprobs"]["content"][0]["top_logprobs"]
            # Exact-match the allowed letters only. The API returns PRE-bias
            # logprobs, so near-duplicates such as " C" sit in the top-k beside
            # "C". An earlier version keyed this dict by token.strip(), which let
            # " C" (p~0) overwrite "C" (p~1) and corrupted about half of all
            # confidences. See the 2026-09-21 correction on the post.
            probs = {t["token"]: math.exp(t["logprob"]) for t in top if t["token"] in by_letter}
            total = sum(probs.values())
            conf = max(probs.values()) / total if total > 0 else None
        except (KeyError, IndexError, TypeError):
            pass

        return (
            by_letter.get(letter),
            usage.get("prompt_tokens", 0),
            usage.get("completion_tokens", 0),
            conf,
        )

    return Backend("logit", run, notes="1 output token, logprobs as confidence")


# ---------------------------------------------------------------------------
# 4. Encoder classifier — the on-prem baseline. No API, no per-token bill,
#    runs on hardware you already own. Cost is amortised GPU time, not tokens.
# ---------------------------------------------------------------------------
def make_encoder(model_path: str) -> Backend:
    try:
        import torch
        from transformers import AutoModelForSequenceClassification, AutoTokenizer
    except ImportError:
        raise SystemExit("encoder backend needs: pip install torch transformers")

    device = "cuda" if torch.cuda.is_available() else "cpu"
    tok = AutoTokenizer.from_pretrained(model_path)
    mdl = AutoModelForSequenceClassification.from_pretrained(model_path).to(device).eval()

    id2label = mdl.config.id2label

    def run(state: str):
        with torch.inference_mode():
            batch = tok(state, return_tensors="pt", truncation=True, max_length=512).to(device)
            logits = mdl(**batch).logits[0]
            probs = torch.softmax(logits, dim=-1)
            idx = int(probs.argmax())
        label = id2label[idx]
        label = label if label in LABELS else None
        n_tok = int(batch["input_ids"].shape[-1])
        return (label, n_tok, 0, float(probs[idx]))

    return Backend("encoder", run, notes=f"local {model_path} on {device}", warmup=3)


# ---------------------------------------------------------------------------
def cost_usd(backend: str, in_tok: int, out_tok: int) -> float:
    p = PRICES.get(backend, {"input": 0.0, "output": 0.0})
    return (in_tok * p["input"] + out_tok * p["output"]) / 1_000_000


def run_backend(be: Backend, cases: list[dict], repeats: int) -> list[Result]:
    results: list[Result] = []

    for _ in range(be.warmup):  # exclude cold start from the latency figures
        try:
            be.run(cases[0]["state"])
        except Exception:
            pass

    for case in cases:
        for _ in range(repeats):
            t0 = time.perf_counter()
            try:
                pred, in_tok, out_tok, conf = be.run(case["state"])
                err = None
            except Exception as exc:
                pred, in_tok, out_tok, conf = None, 0, 0, None
                err = f"{type(exc).__name__}: {exc}"
            dt = (time.perf_counter() - t0) * 1000
            results.append(
                Result(
                    backend=be.name,
                    case_id=case["id"],
                    predicted=pred,
                    expected=case["label"],
                    latency_ms=dt,
                    input_tokens=in_tok,
                    output_tokens=out_tok,
                    confidence=conf,
                    error=err,
                )
            )
            print(f"  {be.name:8} {case['id']:8} {str(pred):10} {dt:7.1f}ms", file=sys.stderr)
    return results


def summarise(results: list[Result], gpu_hourly: float) -> list[dict]:
    rows = []
    by_backend: dict[str, list[Result]] = {}
    for r in results:
        by_backend.setdefault(r.backend, []).append(r)

    for name, rs in by_backend.items():
        ok = [r for r in rs if r.error is None]
        if not ok:
            continue
        lat = sorted(r.latency_ms for r in ok)
        n = len(ok)
        graded = [r for r in ok if r.predicted is not None]
        tok_cost = statistics.mean(
            cost_usd(name, r.input_tokens, r.output_tokens) for r in ok
        )
        # Encoder has no token bill: charge wall-clock GPU time instead.
        if name == "encoder":
            tok_cost = (statistics.mean(lat) / 1000 / 3600) * gpu_hourly

        rows.append(
            {
                "backend": name,
                "n": n,
                "errors": len(rs) - n,
                "unparseable": n - len(graded),
                "accuracy": round(sum(r.correct for r in ok) / n, 4),
                "p50_ms": round(lat[len(lat) // 2], 1),
                "p95_ms": round(lat[min(int(len(lat) * 0.95), len(lat) - 1)], 1),
                "mean_input_tokens": round(statistics.mean(r.input_tokens for r in ok), 1),
                "mean_output_tokens": round(statistics.mean(r.output_tokens for r in ok), 1),
                "usd_per_decision": round(tok_cost, 9),
                "usd_per_million_decisions": round(tok_cost * 1_000_000, 2),
            }
        )
    return sorted(rows, key=lambda r: r["usd_per_decision"])


def main() -> int:
    ap = argparse.ArgumentParser()
    ap.add_argument("--cases", required=True, help="JSONL with id/state/label")
    ap.add_argument("--backends", default="json,logit", help="comma separated")
    ap.add_argument("--model", default="gpt-4o-mini", help="model for json/logit")
    ap.add_argument("--encoder-path", default="./models/route-classifier")
    ap.add_argument("--repeats", type=int, default=3, help="runs per case, for p95")
    ap.add_argument("--gpu-hourly", type=float, default=0.0,
                    help="USD/hr for the encoder's hardware, for amortised cost")
    ap.add_argument("--out", default="decision-layer-bench.csv")
    ap.add_argument("--raw-out", default=None, help="optional per-call CSV")
    ap.add_argument("--labels", default=None,
                    help="JSON {label: description}; overrides the built-in routing labels")
    ap.add_argument("--task", default=None, help="instruction shown to every backend")
    args = ap.parse_args()

    global LABELS, TASK, PRICES
    if args.model not in MODEL_PRICES:
        raise SystemExit(f"no verified price for {args.model!r} — add it to MODEL_PRICES")
    PRICES = {**PRICES, "json": MODEL_PRICES[args.model], "logit": MODEL_PRICES[args.model]}
    if args.labels:
        with open(args.labels) as fh:
            LABELS = json.load(fh)
    if args.task:
        TASK = args.task

    with open(args.cases) as fh:
        cases = [json.loads(line) for line in fh if line.strip()]
    bad = [c["id"] for c in cases if c.get("label") not in LABELS]
    if bad:
        raise SystemExit(f"cases have labels outside LABELS: {bad[:5]}")
    print(f"{len(cases)} cases, {args.repeats} repeats each", file=sys.stderr)

    client = httpx.Client()
    builders = {
        "jev": lambda: make_jev(client),
        "json": lambda: make_json(client, args.model),
        "logit": lambda: make_logit(client, args.model),
        "encoder": lambda: make_encoder(args.encoder_path),
    }

    results: list[Result] = []
    for name in [b.strip() for b in args.backends.split(",") if b.strip()]:
        if name not in builders:
            raise SystemExit(f"unknown backend {name!r}")
        print(f"\n== {name} ==", file=sys.stderr)
        results += run_backend(builders[name](), cases, args.repeats)

    rows = summarise(results, args.gpu_hourly)
    for r in rows:
        if r["backend"] in ("json", "logit"):
            r["backend"] = f"{r['backend']}:{args.model}"
    if not rows:
        raise SystemExit("every call failed — nothing to summarise")

    cols = list(rows[0])
    with open(args.out, "w") as fh:
        fh.write(",".join(cols) + "\n")
        for r in rows:
            fh.write(",".join(str(r[c]) for c in cols) + "\n")

    if args.raw_out:
        with open(args.raw_out, "w") as fh:
            fh.write("backend,case_id,predicted,expected,correct,latency_ms,"
                     "input_tokens,output_tokens,confidence,error\n")
            for r in results:
                fh.write(
                    f"{r.backend},{r.case_id},{r.predicted},{r.expected},{r.correct},"
                    f"{r.latency_ms:.1f},{r.input_tokens},{r.output_tokens},"
                    f"{r.confidence if r.confidence is not None else ''},"
                    f"{r.error or ''}\n"
                )

    print()
    print(f"{'backend':10} {'acc':>7} {'p50':>8} {'p95':>8} {'$/decision':>13} {'$/1M':>10}")
    for r in rows:
        print(f"{r['backend']:10} {r['accuracy']:>7.3f} {r['p50_ms']:>7.1f}ms "
              f"{r['p95_ms']:>7.1f}ms {r['usd_per_decision']:>13.9f} "
              f"{r['usd_per_million_decisions']:>10.2f}")
    print(f"\nwrote {args.out}", file=sys.stderr)
    return 0


if __name__ == "__main__":
    raise SystemExit(main())