"""Probe a deployed endpoint.

SYSTEMONE_URL=$(just url) SYSTEMONE_API_KEY=$(cat .env.key) uv run scripts/probe.py   # growing widths, GPU memory after each
just overhead [app]                                                              # where a request's time goes besides the model
"""

import json
import os
import statistics
import sys
import time

import httpx

URL = os.environ["SYSTEMONE_URL"]
HEADERS = {"authorization": f"Bearer {os.environ['SYSTEMONE_API_KEY']}"}


def choice(width: int) -> dict:
    criteria = {
        f"tool_{i}": f"Tool number {i}: performs operation {i} on the work pool, deployment or flow run resource."
        for i in range(width)
    }
    return {
        "state": {"request": "pause the work pool named production so no new runs start"},
        "model": "kev-latest",
        "questions": {
            "which": {
                "type": "choice",
                "instructions": "Which tool best serves the request?",
                "criteria": criteria,
            }
        },
    }


def main() -> None:
    with httpx.Client(base_url=URL, headers=HEADERS, timeout=600) as client:
        print("at rest", json.dumps(client.get("/health").json()))
        for width in (3, 25, 50, 100, 150):
            t = time.perf_counter()
            r = client.post("/v1/systemone", json=choice(width))
            dt = time.perf_counter() - t
            body = r.json()
            usage = body.get("usage", {}).get("input_tokens")
            print(f"width {width}: {r.status_code} in {dt:.1f}s, input_tokens={usage}")
            print("  ", json.dumps(client.get("/health").json()))
            if r.status_code != 200:
                break


def overhead(rounds: int = 15) -> None:
    with httpx.Client(base_url=URL, headers=HEADERS, timeout=600) as client:
        client.post("/v1/systemone", json=choice(50))
        health, request, model = [], [], []
        for _ in range(rounds):
            t = time.perf_counter()
            client.get("/health")
            health.append((time.perf_counter() - t) * 1000)
            t = time.perf_counter()
            body = client.post("/v1/systemone", json=choice(50)).json()
            request.append((time.perf_counter() - t) * 1000)
            model.append(body["latency_ms"])
    p50 = statistics.median
    print(f"/health round trip   {p50(health):6.0f} ms")
    print(f"50-option request    {p50(request):6.0f} ms")
    print(f"  model (server)     {p50(model):6.0f} ms")
    print(f"  rest in container  {p50(request) - p50(health) - p50(model):6.0f} ms")


if __name__ == "__main__":
    overhead() if sys.argv[1:] == ["overhead"] else main()