#!/usr/bin/env python3
"""Short, same-host ab diagnostic for 1,040-byte HTTP/1 proxy responses.

Requires an existing executable and nginx; keeps the same explicitly supplied
idle-pool limit across variants. Not a substitute for wrk/CPU-isolated acceptance.
"""
from __future__ import annotations
import argparse
import hashlib
import json
from pathlib import Path
import re
import subprocess
import sys
import time

from common import Fixture, http, write_json


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('--binary', action='append', required=True, metavar='LABEL=/abs/path')
    parser.add_argument('--nginx', default='/usr/sbin/nginx', type=Path)
    parser.add_argument('--ab', default='/usr/bin/ab', type=Path)
    parser.add_argument('--out', required=True, type=Path)
    parser.add_argument('--concurrency', default='16,64')
    parser.add_argument('--duration', type=int, default=4)
    parser.add_argument('--repeats', type=int, default=2)
    parser.add_argument('--workers', type=int, default=4)
    parser.add_argument('--max-idle', type=int, default=64)
    args = parser.parse_args()
    binaries = dict(item.split('=', 1) for item in args.binary)
    if len(binaries) != len(args.binary) or not all(Path(x).is_absolute() for x in binaries.values()):
        parser.error('binary labels must be distinct and paths must be absolute')
    if args.duration < 2 or args.repeats < 1:
        parser.error('duration >=2 seconds and repeats >=1 required')
    concurrency = [int(v) for v in args.concurrency.split(',')]
    args.out.mkdir(parents=True, exist_ok=True)
    records = []
    write_json(args.out / 'metadata.json', {
        'simulation': False, 'binary_labels': list(binaries), 'nginx': str(args.nginx),
        'max_idle_for_every_rust_variant': args.max_idle, 'duration_s': args.duration,
        'repeats': args.repeats, 'concurrency': concurrency, 'workers': args.workers,
        'limitations': 'same-host shared CPU and closed-loop ab; not build-level before/after of uncompiled source'
    })
    with Fixture(args.out, workers=args.workers) as fixture:
        nginx = fixture.nginx(args.nginx, frontend=True)
        instances = {name: fixture.rust(Path(binary), name, max_idle=args.max_idle)
                     for name, binary in binaries.items()}
        for name, instance in instances.items():
            instance.login()
            print('Waiting for quiescence:', name, flush=True)
            instance.wait_quiescent(timeout=150, quiet_seconds=3)
        sources = {**{name: r.reverse_port for name, r in instances.items()}, 'nginx': nginx['reverse_port']}
        payload = (fixture.site / 'small.html').read_bytes()
        for name, port in sources.items():
            code, _, body = http(port, '/small.html', headers={'Accept-Encoding': 'identity'})
            if code != 200 or body != payload:
                raise AssertionError(f'{name} precheck failed: {code}, {len(body)} bytes')
        write_json(args.out / 'payload-check.json', {
            'status': 200, 'bytes': len(payload), 'sha256': hashlib.sha256(payload).hexdigest(),
            'all_variants_verified': True
        })
        labels = list(sources)
        for c in concurrency:
            for repeat in range(args.repeats):
                order = labels[repeat % len(labels):] + labels[:repeat % len(labels)]
                if repeat % 2:
                    order.reverse()
                for name in order:
                    port = sources[name]
                    cmd = [str(args.ab), '-k', '-t', str(args.duration), '-n', '1000000', '-c', str(c),
                           '-H', 'Accept-Encoding: identity', f'http://127.0.0.1:{port}/small.html']
                    prefix = f'c{c}-rep{repeat+1}-{name}'
                    subprocess.run([str(args.ab), '-k', '-n', '256', '-c', str(c),
                                   f'http://127.0.0.1:{port}/small.html'],
                                   capture_output=True, timeout=25, check=True)
                    p = subprocess.run(cmd, capture_output=True, text=True, timeout=args.duration + 45)
                    raw = p.stdout + '\n' + p.stderr
                    (args.out / f'{prefix}.txt').write_text(raw)
                    if p.returncode:
                        raise RuntimeError(f'{prefix} ab failed: {p.returncode} {p.stderr}')
                    def field(label):
                        m = re.search(r'^' + re.escape(label) + r':\s*(\d+(?:\.\d+)?)', raw, re.M)
                        return float(m.group(1)) if m else None
                    # ApacheBench tracks both network failures and bad body lengths.
                    if not re.search(r'^Failed requests:\s+0\s*$', raw, re.M):
                        raise RuntimeError(f'{prefix} failed requests: {raw[-1600:]}')
                    if not re.search(r'^Non-2xx responses:\s*0\s*$', raw, re.M) and 'Non-2xx responses:' in raw:
                        raise RuntimeError(f'{prefix} received non-2xx')
                    code, _, body = http(port, '/small.html', headers={'Accept-Encoding': 'identity'})
                    if code != 200 or body != payload:
                        raise AssertionError(f'{prefix} postcheck failed')
                    rps = field('Requests per second')
                    if rps is None:
                        raise RuntimeError(f'{prefix} missing RPS')
                    sample = {'variant': name, 'concurrency': c, 'repeat': repeat+1,
                              'rps': rps, 'completed': field('Complete requests'), 'failed': 0,
                              'expected_sha256': hashlib.sha256(payload).hexdigest(), 'status': 200}
                    records.append(sample)
                    write_json(args.out / 'samples.json', records)
                    print(sample, flush=True)
    write_json(args.out / 'complete.json', {'complete': True, 'samples': len(records)})
    return 0


if __name__ == '__main__':
    sys.exit(main())