#!/usr/bin/env python3
import argparse
import json
import math
import statistics
import sys
import time
import urllib.error
import urllib.request


def percentile(sorted_values, pct):
    if not sorted_values:
        return 0.0
    if len(sorted_values) == 1:
        return float(sorted_values[0])
    rank = (len(sorted_values) - 1) * pct
    low = math.floor(rank)
    high = math.ceil(rank)
    if low == high:
        return float(sorted_values[low])
    weight = rank - low
    return float(sorted_values[low] * (1.0 - weight) + sorted_values[high] * weight)


def wait_for_health(base_url, timeout_s, server_log):
    deadline = time.time() + timeout_s
    last_error = "server did not become healthy"
    while time.time() < deadline:
        try:
            with urllib.request.urlopen(base_url + "/health", timeout=2) as response:
                if response.status == 200:
                    return
        except Exception as exc:  # noqa: BLE001
            last_error = str(exc)
            time.sleep(0.5)

    print("[ERROR] health check failed: {}".format(last_error), file=sys.stderr)
    if server_log:
        try:
            with open(server_log, "r", encoding="utf-8", errors="replace") as fh:
                lines = fh.readlines()[-40:]
            print("[ERROR] tail of server log:", file=sys.stderr)
            for line in lines:
                print(line.rstrip(), file=sys.stderr)
        except OSError as exc:  # noqa: BLE001
            print("[ERROR] failed to read server log: {}".format(exc), file=sys.stderr)
    sys.exit(1)


def build_payload(model_id):
    return {
        "model": model_id,
        "stream": False,
        "temperature": 0.0,
        "top_p": 1.0,
        "max_completion_tokens": 48,
        "messages": [
            {
                "role": "user",
                "content": "Write a compact Rust function that parses a comma-separated list of u32 values and returns a Vec<u32> while ignoring empty segments.",
            }
        ],
    }


def run_once(base_url, model_id):
    body = json.dumps(build_payload(model_id)).encode("utf-8")
    req = urllib.request.Request(
        base_url + "/v1/chat/completions",
        data=body,
        headers={"content-type": "application/json"},
        method="POST",
    )
    started = time.perf_counter()
    with urllib.request.urlopen(req, timeout=180) as response:
        payload = json.loads(response.read().decode("utf-8"))
    elapsed_ms = (time.perf_counter() - started) * 1000.0
    usage = payload.get("usage") or {}
    return elapsed_ms, float(usage.get("tokens_per_second") or 0.0)


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--base-url", required=True)
    parser.add_argument("--model-id", required=True)
    parser.add_argument("--server-log", default="")
    parser.add_argument("--warmups", type=int, default=1)
    parser.add_argument("--samples", type=int, default=7)
    parser.add_argument("--health-timeout", type=float, default=120.0)
    args = parser.parse_args()

    wait_for_health(args.base_url, args.health_timeout, args.server_log)

    for _ in range(args.warmups):
        run_once(args.base_url, args.model_id)

    latencies = []
    tps_values = []
    for _ in range(args.samples):
        latency_ms, tps = run_once(args.base_url, args.model_id)
        latencies.append(latency_ms)
        tps_values.append(tps)

    latencies_sorted = sorted(latencies)
    p50_ms = percentile(latencies_sorted, 0.50)
    p95_ms = percentile(latencies_sorted, 0.95)
    mean_ms = statistics.fmean(latencies)
    completion_tps_mean = statistics.fmean(tps_values) if tps_values else 0.0

    print("METRIC p95_ms={:.3f}".format(p95_ms))
    print("METRIC p50_ms={:.3f}".format(p50_ms))
    print("METRIC mean_ms={:.3f}".format(mean_ms))
    print("METRIC completion_tps_mean={:.3f}".format(completion_tps_mean))
    print("METRIC success_count={}".format(len(latencies)))


if __name__ == "__main__":
    main()