#!/usr/bin/env python3
"""
throughput.py - a minimal, dependency-free concurrent HTTP load generator.

This repo intentionally does not shell out to wrk/hey/ab (neither was available in the
sandbox this repo was built in, and a stdlib script means the load generator itself needs
no separate install or version callout). Each worker thread keeps one persistent
HTTP/1.1 connection (http.client, keep-alive) and fires GET requests back-to-back for a
fixed wall-clock duration; the main thread joins all workers and aggregates.

Usage:
    throughput.py <label> <host> <port> <path> <duration_s> <concurrency>

Prints one JSON line: requests, errors, requests_per_sec, latency_ms mean/p50/p95/p99.
"""
import http.client
import json
import statistics
import sys
import threading
import time


def worker(host, port, path, stop_at, latencies, counters, idx):
    conn = http.client.HTTPConnection(host, port, timeout=5)
    local_ok = 0
    local_err = 0
    local_latencies = []
    while time.monotonic() < stop_at:
        t0 = time.monotonic()
        try:
            conn.request("GET", path)
            resp = conn.getresponse()
            resp.read()
            if resp.status == 200:
                local_ok += 1
            else:
                local_err += 1
        except Exception:
            local_err += 1
            try:
                conn.close()
            except Exception:
                pass
            conn = http.client.HTTPConnection(host, port, timeout=5)
        local_latencies.append((time.monotonic() - t0) * 1000.0)
    counters[idx] = (local_ok, local_err)
    latencies[idx] = local_latencies
    try:
        conn.close()
    except Exception:
        pass


def main():
    label, host, port, path, duration_s, concurrency = sys.argv[1:7]
    port = int(port)
    duration_s = float(duration_s)
    concurrency = int(concurrency)

    counters = [None] * concurrency
    latencies = [None] * concurrency
    stop_at = time.monotonic() + duration_s

    threads = [
        threading.Thread(target=worker, args=(host, port, path, stop_at, latencies, counters, i))
        for i in range(concurrency)
    ]
    wall_start = time.monotonic()
    for t in threads:
        t.start()
    for t in threads:
        t.join()
    wall_elapsed = time.monotonic() - wall_start

    total_ok = sum(c[0] for c in counters)
    total_err = sum(c[1] for c in counters)
    all_latencies = [x for sub in latencies for x in sub]
    all_latencies.sort()

    def pct(p):
        if not all_latencies:
            return None
        k = min(len(all_latencies) - 1, int(len(all_latencies) * p))
        return round(all_latencies[k], 2)

    result = {
        "label": label,
        "path": path,
        "concurrency": concurrency,
        "duration_s": round(wall_elapsed, 2),
        "requests_ok": total_ok,
        "requests_err": total_err,
        "requests_per_sec": round(total_ok / wall_elapsed, 1) if wall_elapsed > 0 else None,
        "latency_ms_mean": round(statistics.mean(all_latencies), 2) if all_latencies else None,
        "latency_ms_p50": pct(0.50),
        "latency_ms_p95": pct(0.95),
        "latency_ms_p99": pct(0.99),
    }
    print(json.dumps(result, indent=2))


if __name__ == "__main__":
    main()