"""
stress_test.py — measure the ceilings that actually matter for this backend.

The GPU is NOT this server's bottleneck: it runs on external ZeroGPU Spaces and is
billed per-user, so you can't meaningfully stress it from one account. The two real
questions are:

  1. WRITE ceiling — how fast can we commit to the dataset before HF starts
     returning 429s? All discoveries funnel through one write token, so this is the
     central limit. (Probe against a THROWAWAY dataset.)
  2. READ path — the frontend reads straight from the dataset CDN; confirm those
     /resolve/ URLs are fast and CORS-open (they shouldn't touch the server at all).

An optional --live probe hits the real /scan endpoint concurrently; it consumes
your GPU quota and creates real monsters, so it's off by default.

Examples:
    # 1) Find the commit-rate ceiling on a scratch dataset you own:
    HF_TOKEN=hf_xxx python stress_test.py commits --repo <you>/piclets-stress --n 60

    # 2) Check the read/CDN path for your real dataset:
    python stress_test.py reads --repo Fraser/piclets --n 200 --concurrency 20

    # 3) (optional) Live end-to-end concurrency against the deployed Space:
    python stress_test.py live --space <you>/piclets-server --image test.jpg \
        --token hf_xxx --n 5 --concurrency 3
"""

import io
import os
import sys
import time
import json
import argparse
import statistics
import concurrent.futures as cf

import requests


# --------------------------------------------------------------------------
# 1) WRITE ceiling: burst commits to a throwaway dataset, watch for 429s
# --------------------------------------------------------------------------
def probe_commits(repo: str, n: int):
    from huggingface_hub import HfApi, CommitOperationAdd
    from huggingface_hub.utils import HfHubHTTPError

    token = os.getenv("HF_TOKEN")
    if not token:
        sys.exit("Set HF_TOKEN (write access to the scratch repo).")
    api = HfApi(token=token)
    api.create_repo(repo, repo_type="dataset", exist_ok=True, private=True)
    print(f"[commits] probing {repo} with {n} sequential commits "
          f"(huggingface_hub auto-retries 429; watch for slpowdowns)...\n")

    latencies, throttles, errors = [], 0, 0
    t0 = time.time()
    for i in range(n):
        blob = io.BytesIO(json.dumps({"i": i, "t": time.time()}).encode())
        op = CommitOperationAdd(path_in_repo=f"stress/{i}.json", path_or_fileobj=blob)
        start = time.time()
        try:
            api.create_commit(repo_id=repo, repo_type="dataset",
                              operations=[op], commit_message=f"stress {i}")
            dt = time.time() - start
            latencies.append(dt)
            # A commit that took much longer than the median usually means the SDK
            # slept off a 429 for us.
            flag = " <- slow (likely throttled + retried)" if dt > 3 else ""
            print(f"  commit {i:>3}: {dt:5.2f}s{flag}")
        except HfHubHTTPError as exc:
            code = getattr(getattr(exc, "response", None), "status_code", "?")
            if code == 429:
                throttles += 1
                print(f"  commit {i:>3}: 429 rate-limited")
            else:
                errors += 1
                print(f"  commit {i:>3}: HTTP {code}")
        except Exception as exc:
            errors += 1
            print(f"  commit {i:>3}: {type(exc).__name__}: {exc}")
    total = time.time() - t0

    print("\n[commits] summary")
    print(f"  committed:        {len(latencies)}/{n}")
    print(f"  429s (uncaught):  {throttles}")
    print(f"  other errors:     {errors}")
    if latencies:
        print(f"  median latency:   {statistics.median(latencies):.2f}s")
        print(f"  p90 latency:      {sorted(latencies)[int(len(latencies)*0.9)-1]:.2f}s")
        print(f"  sustained rate:   {len(latencies)/total:.2f} commits/s over {total:.1f}s")
    print("\n  Interpretation: the rate where latency climbs / retries kick in is your\n"
          "  practical write ceiling. Since one token serves ALL discoveries, that's the\n"
          "  number that bounds global new-monster throughput. Remember to delete this\n"
          "  scratch dataset afterwards.")


# --------------------------------------------------------------------------
# 2) READ path: hammer a /resolve/ URL, check latency + CORS
# --------------------------------------------------------------------------
def probe_reads(repo: str, n: int, concurrency: int):
    url = f"https://huggingface.co/datasets/{repo}/resolve/main/index/stats.json"
    print(f"[reads] GET {url}\n        {n} requests, {concurrency} concurrent\n")

    def one(_):
        start = time.time()
        r = requests.get(url, timeout=30)
        return time.time() - start, r.status_code, r.headers.get("access-control-allow-origin")

    lat, codes, cors_seen = [], {}, set()
    with cf.ThreadPoolExecutor(max_workers=concurrency) as ex:
        for dt, code, cors in ex.map(one, range(n)):
            lat.append(dt)
            codes[code] = codes.get(code, 0) + 1
            if cors:
                cors_seen.add(cors)

    print("[reads] summary")
    print(f"  status codes:     {codes}")
    print(f"  median latency:   {statistics.median(lat)*1000:.0f} ms")
    print(f"  p90 latency:      {sorted(lat)[int(len(lat)*0.9)-1]*1000:.0f} ms")
    print(f"  CORS header seen: {cors_seen or 'NONE (check if the browser can read this cross-origin)'}")
    print("\n  Interpretation: reads are per-client and CDN-served, so this should stay fast\n"
          "  under load and never involve the server. A '*' or your origin in CORS means\n"
          "  the frontend can fetch it directly.")


# --------------------------------------------------------------------------
# 3) LIVE end-to-end (optional; spends GPU quota, creates real monsters)
# --------------------------------------------------------------------------
def probe_live(space: str, image: str, token: str, n: int, concurrency: int):
    from gradio_client import Client, handle_file

    if not os.path.exists(image):
        sys.exit(f"Image not found: {image}")
    print(f"[live] WARNING: this spends your ZeroGPU quota and writes real monsters.\n"
          f"       {n} scans, {concurrency} concurrent, against {space}\n")

    def one(_):
        start = time.time()
        try:
            client = Client(space, hf_token=token)
            res = client.predict(handle_file(image), token, api_name="/scan")
            return time.time() - start, (res.get("status") if isinstance(res, dict) else "?"), None
        except Exception as exc:
            return time.time() - start, "error", str(exc)[:120]

    lat, outcomes = [], {}
    with cf.ThreadPoolExecutor(max_workers=concurrency) as ex:
        for dt, status, err in ex.map(one, range(n)):
            lat.append(dt)
            outcomes[status] = outcomes.get(status, 0) + 1
            if err:
                print(f"  error: {err}")

    print("\n[live] summary")
    print(f"  outcomes:        {outcomes}")
    print(f"  median latency:  {statistics.median(lat):.1f}s")
    print(f"  p90 latency:     {sorted(lat)[int(len(lat)*0.9)-1]:.1f}s")


def main():
    p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    sub = p.add_subparsers(dest="cmd", required=True)

    c = sub.add_parser("commits", help="probe the dataset commit-rate ceiling (throwaway repo)")
    c.add_argument("--repo", required=True)
    c.add_argument("--n", type=int, default=60)

    r = sub.add_parser("reads", help="probe the CDN read path")
    r.add_argument("--repo", required=True)
    r.add_argument("--n", type=int, default=200)
    r.add_argument("--concurrency", type=int, default=20)

    l = sub.add_parser("live", help="optional live end-to-end scan probe")
    l.add_argument("--space", required=True)
    l.add_argument("--image", required=True)
    l.add_argument("--token", required=True)
    l.add_argument("--n", type=int, default=5)
    l.add_argument("--concurrency", type=int, default=3)

    args = p.parse_args()
    if args.cmd == "commits":
        probe_commits(args.repo, args.n)
    elif args.cmd == "reads":
        probe_reads(args.repo, args.n, args.concurrency)
    elif args.cmd == "live":
        probe_live(args.space, args.image, args.token, args.n, args.concurrency)


if __name__ == "__main__":
    main()