#!/usr/bin/env python3
"""Finite live HTTP concurrency correctness stress; no throughput/latency claim."""
from __future__ import annotations

import argparse
from concurrent.futures import ThreadPoolExecutor
import hashlib
import json
from pathlib import Path
import threading
import time
import traceback

from common import Fixture, http, seal_results, sha256, write_json
from resources import process_stats


PAYLOADS = ["small.html", "large.bin", "huge.bin"]


def active_backends(rust):
    status, _, body = rust.admin("/api/status")
    if status != 200:
        return {"error": f"HTTP {status}"}
    values = []

    def walk(value):
        if isinstance(value, dict):
            if "id" in value and "active" in value and "status" in value and "requests" in value:
                values.append({key: value[key] for key in ("id", "active", "status", "requests")})
            for child in value.values():
                walk(child)
        elif isinstance(value, list):
            for child in value:
                walk(child)
    walk(json.loads(body))
    return values


def batch(rust, port, name, expected, concurrency):
    count = max(32, concurrency * 2)
    expected_sha = hashlib.sha256(expected).hexdigest()
    ready = threading.Barrier(concurrency, timeout=15)
    lock = threading.Lock()
    running = peak = 0

    def fetch(index):
        nonlocal running, peak
        with lock:
            running += 1
            peak = max(peak, running)
        try:
            if index < concurrency:
                ready.wait()
            status, headers, body = http(port, "/" + name, timeout=15)
            digest = hashlib.sha256(body).hexdigest()
            return {"index": index, "status": status, "bytes": len(body), "sha256": digest,
                    "passed": status == 200 and len(body) == len(expected) and digest == expected_sha}
        except Exception as error:
            return {"index": index, "passed": False, "error": str(error)}
        finally:
            with lock:
                running -= 1

    before = process_stats(rust.process.pid)
    started = time.monotonic()
    with ThreadPoolExecutor(max_workers=concurrency) as executor:
        responses = list(executor.map(fetch, range(count)))
    elapsed = time.monotonic() - started
    # Completion means every client consumed or failed its full response. Give
    # asynchronous server bookkeeping a bounded interval to release its leases.
    for _ in range(40):
        active = active_backends(rust)
        if isinstance(active, list) and active and all(row["active"] == 0 for row in active):
            break
        time.sleep(.05)
    after = process_stats(rust.process.pid)
    return {"configured_concurrent_clients": concurrency, "peak_running_client_tasks": peak,
            "client_tasks_include_barrier": True, "requests": count, "elapsed_seconds": elapsed,
            "expected_bytes": len(expected), "expected_sha256": expected_sha,
            "response_failures": sum(not response["passed"] for response in responses),
            "passed": all(response["passed"] for response in responses)
                      and isinstance(active, list) and bool(active) and all(row["active"] == 0 for row in active),
            "responses": responses, "active_backends_after_drain": active,
            "server_before": before, "server_after_drain": after}


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--binary", action="append", required=True, metavar="LABEL=/ABS/PATH")
    parser.add_argument("--nginx", type=Path, required=True)
    parser.add_argument("--out", type=Path, required=True)
    args = parser.parse_args()
    if args.out.exists() and any(args.out.iterdir()):
        parser.error("--out must be new or empty")
    reports = {}
    for specification in args.binary:
        label, binary = specification.split("=", 1)
        cases, error = [], None
        try:
            with Fixture(args.out / label, workers=4, origin_workers=2) as fixture:
                fixture.nginx(args.nginx, frontend=False)
                rust = fixture.rust(Path(binary), label)
                rust.login()
                for source, port in [("proxy", rust.reverse_port), ("static", rust.static_port)]:
                    for name in PAYLOADS:
                        expected = (fixture.site / name).read_bytes()
                        for concurrency in [1, 16, 128]:
                            result = batch(rust, port, name, expected, concurrency)
                            case = {"case": f"concurrency.{source}.{name}.c{concurrency}", **result}
                            cases.append(case)
                            write_json(args.out / label / "cases" / f"{source}-{name}-c{concurrency}.json", case)
                            print(f"{label} {'PASS' if case['passed'] else 'FAIL'} {case['case']}", flush=True)
                if rust.process.poll() is not None:
                    raise RuntimeError("RustProxy exited during concurrency regression")
        except Exception as exc:
            error = {"error": str(exc), "traceback": traceback.format_exc()}
            write_json(args.out / label / "failure.json", error)
        report = {"kind": "live HTTP/1.1 finite concurrency functional stress", "simulation": False,
                  "binary_sha256": sha256(Path(binary)), "throughput_measurement": False,
                  "passed_cases": sum(case["passed"] for case in cases),
                  "failed_cases": sum(not case["passed"] for case in cases),
                  "requests": sum(case["requests"] for case in cases),
                  "response_failures": sum(case["response_failures"] for case in cases),
                  "completed": error is None and len(cases) == 18,
                  "cases": [{key: case[key] for key in ("case", "passed", "requests", "response_failures")} for case in cases],
                  "limitations": ["configured client concurrency is not proof of simultaneous active server exchanges",
                                  "Python client is not a throughput or latency benchmark",
                                  "FD/RSS are before/after snapshots, not peaks or a long-term leak test"]}
        reports[label] = report
        write_json(args.out / label / "report.json", report)
    write_json(args.out / "comparison.json", reports)
    seal_results(args.out)
    return 0 if all(report["completed"] and not report["failed_cases"] for report in reports.values()) else 1


if __name__ == "__main__":
    raise SystemExit(main())