"""Latency / concurrency check for the served classifier.

  python scripts/load_test.py --url http://localhost:8000 --clients 8 --requests 8 --batch 1
Reports per-request latency (single input), and throughput with N concurrent clients each
sending `batch` inputs per request. Also asserts every response is label-consistent with a
sequential baseline (state isolation between requests).
"""
import argparse
import json
import statistics
import threading
import time
import urllib.request

TEMPLATE = "Premise: {premise}\nHypothesis: {hypothesis}"
PAIRS = [
    ("A man is playing a guitar.", "Someone is making music."),
    ("A man is playing a guitar.", "A woman is sleeping."),
    ("Two dogs are running through a field.", "There are animals outdoors."),
    ("It is raining heavily outside.", "The weather is dry."),
    ("She bought three apples and two oranges at the market.", "She purchased fruit."),
    ("The meeting was postponed because the CEO was ill.", "The CEO was not well."),
    ("The Eiffel Tower is located in Paris, France, and was completed in 1889.", "The Eiffel Tower is in Germany."),
    ("Which gas do plants absorb during photosynthesis?", "The correct answer is: carbon dioxide"),
]


def post(url, payload, timeout=600):
    req = urllib.request.Request(url, data=json.dumps(payload).encode(), headers={"Content-Type": "application/json"})
    with urllib.request.urlopen(req, timeout=timeout) as r:
        return json.loads(r.read())


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--url", default="http://localhost:8000")
    ap.add_argument("--model", default="AlexWortega/openjev")
    ap.add_argument("--clients", type=int, default=8)
    ap.add_argument("--requests", type=int, default=8, help="requests per client")
    ap.add_argument("--batch", type=int, default=1, help="inputs per request")
    args = ap.parse_args()
    texts = [TEMPLATE.format(premise=p, hypothesis=h) for p, h in PAIRS]

    # Sequential baseline: single-input requests.
    base = []
    lat = []
    for t in texts:
        t0 = time.time()
        out = post(f"{args.url}/classify", {"model": args.model, "input": [t]})
        lat.append(time.time() - t0)
        base.append(out["data"][0]["probs"])
    print(f"single-input latency: median {statistics.median(lat)*1e3:.0f} ms, min {min(lat)*1e3:.0f} ms, max {max(lat)*1e3:.0f} ms")

    # One request with all inputs (batched prefill step).
    t0 = time.time()
    out = post(f"{args.url}/classify", {"model": args.model, "input": texts})
    dt = time.time() - t0
    print(f"one request x {len(texts)} inputs: {dt*1e3:.0f} ms total, {dt/len(texts)*1e3:.0f} ms/input")
    for got, ref in zip(out["data"], base):
        assert max(abs(a - b) for a, b in zip(got["probs"], ref)) < 1e-3, (got["probs"], ref)

    # Concurrent clients.
    results, errors = [], []
    lock = threading.Lock()

    def client(cid):
        for r in range(args.requests):
            idx = [(cid + r + k) % len(texts) for k in range(args.batch)]
            t0 = time.time()
            try:
                out = post(f"{args.url}/classify", {"model": args.model, "input": [texts[i] for i in idx]})
                dt = time.time() - t0
                for k, d in enumerate(out["data"]):
                    ref = base[idx[k]]
                    if max(abs(a - b) for a, b in zip(d["probs"], ref)) >= 1e-3:
                        with lock:
                            errors.append((cid, r, idx[k], d["probs"], ref))
                with lock:
                    results.append(dt)
            except Exception as e:  # noqa: BLE001
                with lock:
                    errors.append((cid, r, repr(e)))

    t0 = time.time()
    threads = [threading.Thread(target=client, args=(c,)) for c in range(args.clients)]
    [t.start() for t in threads]
    [t.join() for t in threads]
    wall = time.time() - t0
    n_inputs = args.clients * args.requests * args.batch
    print(f"{args.clients} clients x {args.requests} req x {args.batch} inputs: {n_inputs} inputs in {wall:.2f}s "
          f"= {n_inputs/wall:.1f} inputs/s; per-request latency median {statistics.median(results)*1e3:.0f} ms "
          f"p95 {sorted(results)[int(0.95*len(results))-1]*1e3:.0f} ms")
    if errors:
        print("ERRORS/MISMATCHES:", errors[:5])
        raise SystemExit(1)
    print("OK: all concurrent results identical to the sequential baseline")


if __name__ == "__main__":
    main()