#!/usr/bin/env python3
"""Record Coral API timings and selected health snapshots before/after a change."""
import argparse, concurrent.futures, datetime, json, pathlib, time, urllib.request

PATHS = ['/health', '/stats', '/simcluster/stats', '/entity-graph',
         '/simcluster/entity-graph', '/state/32', '/diagnostics',
         '/diagnostics/history/views', '/groups/history?limit=10',
         '/simcluster/groups/history?limit=10'] + [
    f'{scope}/history/{kind}?range={span}'
    for scope in ('', '/simcluster') for kind in ('top', 'topics')
    for span in ('hour', 'day', 'week')]

def probe(path):
    start = time.monotonic()
    record = {'path': path, 'at': datetime.datetime.now(datetime.timezone.utc).isoformat()}
    body = None
    try:
        with urllib.request.urlopen('https://coral.fly.dev' + path, timeout=20) as response:
            record.update(status=response.status, ttfb_ms=round((time.monotonic()-start)*1000),
                          server_timing=response.headers.get('Server-Timing'))
            raw = response.read()
            body = json.loads(raw)
            record['bytes'] = len(raw)
            record['shape'] = sorted(body) if isinstance(body, dict) else type(body).__name__
            if isinstance(body, dict):
                for key in ('generated_at', 'source_through'):
                    if key in body: record[key] = body[key]
    except Exception as error:
        record['error'] = str(error)
    record['total_ms'] = round((time.monotonic()-start)*1000)
    return record, body

def main():
    p = argparse.ArgumentParser(description=__doc__)
    p.add_argument('output', type=pathlib.Path)
    p.add_argument('--rounds', type=int, default=6)
    p.add_argument('--interval', type=float, default=10)
    args = p.parse_args()
    records, snapshots = [], {}
    with concurrent.futures.ThreadPoolExecutor(max_workers=4) as pool:
        for i in range(args.rounds):
            for record, body in pool.map(probe, PATHS):
                records.append(record)
                if record['path'] in ('/health', '/stats', '/simcluster/stats', '/diagnostics', '/diagnostics/history/views'):
                    snapshots[record['path']] = body
            if i+1 < args.rounds: time.sleep(args.interval)
    args.output.write_text(json.dumps({'requests':records,'snapshots':snapshots}, indent=2)+'\n')
    print(json.dumps({'requests':len(records),'failures':sum(r.get('status')!=200 for r in records),
                      'max_ms':max(r['total_ms'] for r in records)}))

if __name__ == '__main__': main()