#!/usr/bin/env python3
"""Minimal client for a dual-head bundle served by `vllm serve` (no jev_judge import needed).

Server (one engine, both heads):
  vllm serve autotrust/JEV --served-model-name autotrust/JEV \
      --enable-lora --max-lora-rank 32 --lora-modules jev-decision=<local path>/adapter_vllm \
      --logprobs-mode processed_logprobs --max-model-len 4096

Test:  python3 scripts/vllm_openai_client.py --bundle exports/jev-judge-qwen35-9b-v0.8 --base-url http://localhost:18090
"""

from __future__ import annotations

import argparse
import json
import math
import os

import requests

LETTERS = "ABCDEFGHIJKLMNOP"


class JevClient:
    def __init__(self, base_url: str, bundle_dir: str, decision_model: str = "jev-decision"):
        self.url = base_url.rstrip("/")
        self.model = decision_model
        dh = json.load(open(os.path.join(bundle_dir, "adapter_vllm", "decision_head.json")))
        self.bias, self.ids, self.ranges = dh["bias"], dh["verbalizer_ids"], dh["slots"]["ranges"]
        self.temp = json.load(open(os.path.join(bundle_dir, "calibration.json")))["per_kind"]

    @staticmethod
    def prompt(kind: str, state: str, question: str, options: list[str]) -> str:
        lines = options if kind != "choice" else [f"{LETTERS[i]}) {o}" for i, o in enumerate(options)]
        return f"[kind] {kind}\n[state] {state}\n[question] {question}\n[options]\n" + "\n".join(lines) + "\n[decision]:"

    def decide(self, kind: str, state: str, question: str, options: list[str], session: requests.Session | None = None) -> dict:
        if kind == "noul":
            options = ["false", "true"]
        elif kind == "score":
            options = ["0", "1", "2", "3", "4", "5"]
        s = self.ranges[kind][0]
        n = len(options)
        allowed = self.ids[s : s + n]
        body = {"model": self.model, "prompt": self.prompt(kind, state, question, options), "max_tokens": 1,
                "temperature": 1.0, "logprobs": n, "allowed_token_ids": allowed, "add_special_tokens": False,
                "return_tokens_as_token_ids": True}
        r = (session or requests).post(f"{self.url}/v1/completions", json=body, timeout=60)
        r.raise_for_status()
        top = r.json()["choices"][0]["logprobs"]["top_logprobs"][0]  # {"token_id:123": logprob}
        lp = {int(k.split(":")[1]): v for k, v in top.items()}
        z = [(lp.get(t, -math.inf) + self.bias[s + i]) / self.temp[kind] for i, t in enumerate(allowed)]
        m = max(z)
        e = [math.exp(x - m) for x in z]
        p = [x / sum(e) for x in e]
        return dict(zip(options, p))

    def chat(self, messages: list[dict], model: str, **kw) -> str:
        r = requests.post(f"{self.url}/v1/chat/completions", json={"model": model, "messages": messages, **kw}, timeout=300)
        r.raise_for_status()
        return r.json()["choices"][0]["message"]["content"]


def main() -> None:
    import sys
    import time
    from concurrent.futures import ThreadPoolExecutor

    import numpy as np

    sys.path.insert(0, os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "src"))
    from jev_judge.data import load_split
    from jev_judge.template import SLOT_RANGES

    ap = argparse.ArgumentParser()
    ap.add_argument("--bundle", required=True)
    ap.add_argument("--base-url", default="http://localhost:18090")
    ap.add_argument("--served-model", default="autotrust/JEV")
    ap.add_argument("--hf-ref", default=None)
    ap.add_argument("--n", type=int, default=2000)
    ap.add_argument("--concurrency", type=int, default=64)
    ap.add_argument("--out", default=None)
    args = ap.parse_args()
    c = JevClient(args.base_url, args.bundle)

    print("AR   :", c.chat([{"role": "user", "content": "In one sentence, what is safety stock?"}], args.served_model,
                          max_tokens=60, temperature=0.0).replace("\n", " ")[:160])
    print("noul :", c.decide("noul", "Customer says the parcel arrived damaged and wants their money back.",
                            "Is the customer asking for a refund?", []))
    print("choice:", c.decide("choice", "SKU AX-330 stock at 8% of safety level; supplier late twice this quarter.",
                             "Supplier response for this scenario.", ["issue_warning", "renegotiate", "dual_source", "maintain"]))

    df = load_split("data", "test_set_30k").head(args.n)
    rows = list(zip(df["kind"], df["state"], df["question"], df["options"]))
    sess = requests.Session()
    adapter = requests.adapters.HTTPAdapter(pool_connections=args.concurrency, pool_maxsize=args.concurrency)
    sess.mount("http://", adapter)
    for k, s_, q, o in rows[:32]:
        c.decide(k, s_, q, list(o), sess)  # warm-up
    t0 = time.time()
    lat = []

    def one(row):
        t = time.time()
        out = c.decide(row[0], row[1], row[2], list(row[3]), sess)
        lat.append((time.time() - t) * 1000)
        return out

    with ThreadPoolExecutor(args.concurrency) as ex:
        outs = list(ex.map(one, rows))
    dt = time.time() - t0
    res = {"n": len(rows), "concurrency": args.concurrency, "req_per_s": len(rows) / dt, "p50_ms": float(np.median(lat)),
           "p90_ms": float(np.percentile(lat, 90))}
    print(f"decisions over HTTP: {len(rows)} requests, {args.concurrency} concurrent -> {res['req_per_s']:.0f} req/s, "
          f"p50 {res['p50_ms']:.0f} ms, p90 {res['p90_ms']:.0f} ms")
    if args.hf_ref:
        ref = np.load(args.hf_ref)["probs"][: len(rows)]
        d = []
        for (k, _, _, o), out, rr in zip(rows, outs, ref):
            s0 = SLOT_RANGES[k][0]
            d.extend(abs(v - rr[s0 + i]) for i, v in enumerate(out.values()))
        res["vs_pytorch_max_abs_dp"] = float(max(d))
        res["vs_pytorch_mean_abs_dp"] = float(np.mean(d))
        print(f"vs PyTorch path: max|Δp| {max(d):.4f}, mean|Δp| {np.mean(d):.5f}")
    if args.out:
        json.dump(res, open(args.out, "w"), indent=1)


if __name__ == "__main__":
    main()