#!/usr/bin/env python3
"""Measure where serial database work changes the best local app host."""

import argparse
import csv
import json
import random
import statistics
import time
from pathlib import Path
from urllib.error import HTTPError
from urllib.request import urlopen


COUNTS = (0, 1, 3, 6, 12, 24, 48)
SITES = ("A", "B")


def request(base, path):
    started = time.perf_counter_ns()
    try:
        with urlopen(base + path, timeout=10) as response:
            body = response.read()
            status = response.status
    except HTTPError as error:
        body, status = error.read(), error.code
    elapsed_ms = (time.perf_counter_ns() - started) / 1_000_000
    return elapsed_ms, status, json.loads(body) if body else {}


def measure(base, path, count):
    samples = []
    db_samples = []
    for _ in range(count):
        elapsed_ms, status, result = request(base, path)
        if status not in (200, 204):
            raise RuntimeError(f"{path} returned HTTP {status}: {result}")
        samples.append(elapsed_ms)
        if "database_ms" in result:
            db_samples.append(float(result["database_ms"]))
    return samples, db_samples


def bootstrap_delta(a, b, rng, iterations=2000):
    deltas = []
    for _ in range(iterations):
        a_sample = [a[rng.randrange(len(a))] for _ in a]
        b_sample = [b[rng.randrange(len(b))] for _ in b]
        deltas.append(statistics.median(a_sample) - statistics.median(b_sample))
    deltas.sort()
    return deltas[int(iterations * 0.025)], deltas[int(iterations * 0.975)]


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--compute-a-url", required=True)
    parser.add_argument("--compute-b-url", required=True)
    parser.add_argument("--output", type=Path, required=True)
    parser.add_argument("--samples", type=int, default=300)
    parser.add_argument("--warmup", type=int, default=20)
    parser.add_argument("--counts", type=int, nargs="+", default=COUNTS)
    parser.add_argument("--reuse-calibration", type=Path)
    parser.add_argument("--seed", type=int, default=20260926)
    args = parser.parse_args()
    if args.samples < 50 or args.warmup < 1:
        parser.error("samples must be at least 50 and warmup must be positive")

    bases = {"A": args.compute_a_url, "B": args.compute_b_url}
    rng = random.Random(args.seed)
    args.output.mkdir(parents=True, exist_ok=True)
    if args.reuse_calibration:
        calibration = json.loads(args.reuse_calibration.read_text())
    else:
        calibration = {}
        for site, base in bases.items():
            for path in ("/health", "/serial?count=0", "/db-probe"):
                measure(base, path, args.warmup)
            health, _ = measure(base, "/health", args.samples)
            baseline, _ = measure(base, "/serial?count=0", args.samples)
            _, db = measure(base, "/db-probe", args.samples)
            calibration[site] = {
                "health_p50_ms": statistics.median(health),
                "baseline_p50_ms": statistics.median(baseline),
                "db_rtt_p50_ms": statistics.median(db),
            }

    counts = tuple(args.counts)
    if any(count < 0 or count > 64 for count in counts):
        parser.error("each count must be between 0 and 64")
    for count in counts:
        for site, base in bases.items():
            measure(base, f"/serial?count={count}", args.warmup)

    trials = [(count, site, repeat) for count in counts for site in SITES
              for repeat in range(args.samples)]
    rng.shuffle(trials)
    observations = {(count, site): [] for count in counts for site in SITES}
    raw_rows = []
    for count, site, repeat in trials:
        path = f"/serial?count={count}"
        elapsed_ms, status, result = request(bases[site], path)
        if status != 200 or result.get("sequential_queries") != count:
            raise RuntimeError(f"invalid response from site {site} for N={count}")
        observations[(count, site)].append(elapsed_ms)
        raw_rows.append({
            "n": count,
            "site": site,
            "repeat": repeat,
            "latency_ms": f"{elapsed_ms:.6f}",
            "database_ms": f"{float(result['database_ms']):.6f}",
        })

    with (args.output / "raw.csv").open("w", newline="") as output:
        writer = csv.DictWriter(output, fieldnames=raw_rows[0].keys())
        writer.writeheader()
        writer.writerows(raw_rows)
    (args.output / "calibration.json").write_text(
        json.dumps(calibration, indent=2) + "\n"
    )

    rows = []
    for count in counts:
        a, b = observations[(count, "A")], observations[(count, "B")]
        delta = statistics.median(a) - statistics.median(b)
        low, high = bootstrap_delta(a, b, rng)
        health_cost = {
            site: calibration[site]["health_p50_ms"]
            + count * calibration[site]["db_rtt_p50_ms"]
            for site in SITES
        }
        profile_cost = {
            site: calibration[site]["baseline_p50_ms"]
            + count * calibration[site]["db_rtt_p50_ms"]
            for site in SITES
        }
        rows.append({
            "n": count,
            "a_p50_ms": f"{statistics.median(a):.6f}",
            "b_p50_ms": f"{statistics.median(b):.6f}",
            "a_minus_b_ms": f"{delta:.6f}",
            "ci95_low_ms": f"{low:.6f}",
            "ci95_high_ms": f"{high:.6f}",
            "user_first": min(SITES, key=lambda site: calibration[site]["health_p50_ms"]),
            "network_model": min(SITES, key=health_cost.get),
            "profiled_model": min(SITES, key=profile_cost.get),
            "observed_best": "B" if delta > 0 else "A",
        })
    with (args.output / "summary.csv").open("w", newline="") as output:
        writer = csv.DictWriter(output, fieldnames=rows[0].keys())
        writer.writeheader()
        writer.writerows(rows)
    print(json.dumps({"calibration": calibration, "summary": rows}, indent=2))


if __name__ == "__main__":
    main()