#!/usr/bin/env python3
"""Join game traffic samples with process telemetry; preserve individual rounds."""
import argparse
import collections
import json
import pathlib
import statistics


def summarize(traffic, resources):
    samples = []
    for source in traffic['samples']:
        row = dict(source)
        window = [r for r in resources if row['started_epoch'] <= r['epoch'] <= row['finished_epoch']]
        if len(window) < 2:
            raise ValueError(f"Missing telemetry: {row['case']} round {row['round']}")
        elapsed = window[-1]['epoch'] - window[0]['epoch']
        row['resource_window_seconds'] = elapsed
        row['resources'] = {}
        for label in ('server', 'backend'):
            row['resources'][label] = {
                'cpu_cores': (window[-1][label]['cpu_seconds'] - window[0][label]['cpu_seconds']) / elapsed,
                'peak_pss_mib': max(v[label]['pss_mib'] for v in window),
                'peak_rss_mib': max(v[label]['rss_mib'] for v in window),
                'peak_fds': max(v[label]['fds'] for v in window),
            }
        row['host_peak_used_mib'] = max(v['host_used_mib'] for v in window)
        samples.append(row)
    grouped = collections.defaultdict(list)
    for row in samples:
        grouped[row['case']].append(row)
    result = []
    for case, rows in grouped.items():
        def median(get):
            values = [get(r) for r in rows if get(r) is not None]
            return statistics.median(values) if values else None
        errors = collections.Counter()
        for r in rows:
            errors.update(r['errors'])
        summary = {
            'case': case, 'rounds': len(rows), 'attempted_per_round': rows[0]['attempted_connections'],
            'connections_min': min(r['established_connections'] for r in rows),
            'connections_max': max(r['established_connections'] for r in rows),
            'median_delivered_per_second': median(lambda r: r['delivered_per_second']),
            'median_p99_ms': median(lambda r: r['latency_ms']['p99']),
            'median_p50_ms': median(lambda r: r['latency_ms']['p50']),
            'p99_min_ms': min((r['latency_ms']['p99'] for r in rows if r['latency_ms']['p99'] is not None), default=None),
            'p99_max_ms': max((r['latency_ms']['p99'] for r in rows if r['latency_ms']['p99'] is not None), default=None),
            'median_handshake_p99_ms': median(lambda r: r['handshake_ms']['p99']),
            'median_client_lag_p99_ms': median(lambda r: r['client_schedule_lag_ms']['p99']),
            'median_asset_p99_ms': median(lambda r: r['asset_latency_ms']['p99']),
            'asset_requests_total': sum(r['asset_requests'] for r in rows),
            'asset_rounds_completed': sum(r['asset_requests'] > 0 for r in rows),
            'verified_total': sum(r['verified'] for r in rows),
            'errors_total': dict(errors),
            'median_client_cpu_cores': median(lambda r: r['client_cpu_cores']),
            'peak_host_used_mib': max(r['host_peak_used_mib'] for r in rows),
        }
        for label in ('server', 'backend'):
            summary[label] = {
                'median_cpu_cores': median(lambda r: r['resources'][label]['cpu_cores']),
                'peak_pss_mib': max(r['resources'][label]['peak_pss_mib'] for r in rows),
                'peak_rss_mib': max(r['resources'][label]['peak_rss_mib'] for r in rows),
            }
        result.append(summary)
    return {'summary': result, 'samples': samples}


if __name__ == '__main__':
    p = argparse.ArgumentParser(description=__doc__)
    p.add_argument('directory', type=pathlib.Path)
    p.add_argument('--output', type=pathlib.Path, required=True)
    a = p.parse_args()
    result = {}
    for label in ('server-before', 'server-after', 'nano-before', 'nano-after', 'server-sustained'):
        raw = json.loads((a.directory / (label + '.json')).read_text())
        monitor = 'server-after' if label == 'server-sustained' else label
        resources = [json.loads(line) for line in (a.directory / (monitor + '-resources.jsonl')).read_text().splitlines()]
        result[label] = summarize(raw, resources)
    a.output.write_text(json.dumps(result, indent=2) + '\n')