#!/usr/bin/env python3
"""Measure xz/zstd on reproducible, generated data; never reads user data."""
import argparse
import gzip
import hashlib
import json
import os
from pathlib import Path
import platform
import random
import statistics
import subprocess
import tempfile
import time
import zlib


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('--output', required=True, type=Path)
    parser.add_argument('--repeats', default=3, type=int)
    args = parser.parse_args()
    if args.repeats < 1:
        parser.error('--repeats must be positive')
    env = os.environ.copy()
    for key in ('XZ_OPT', 'XZ_DEFAULTS', 'ZSTD_CLEVEL', 'ZSTD_NBTHREADS'):
        env.pop(key, None)
    configs = [('xz-3', ['xz', '-T1', '-3', '-c'], 'xz'),
               ('xz-6', ['xz', '-T1', '-6', '-c'], 'xz'),
               ('zstd-3', ['zstd', '--single-thread', '-3', '-c'], 'zstd'),
               ('zstd-9', ['zstd', '--single-thread', '-9', '-c'], 'zstd')]
    cpu = next((line.split(':', 1)[1].strip()
                for line in Path('/proc/cpuinfo').read_text().splitlines()
                if line.startswith('model name')), platform.machine())
    report = {'schema_version': 1, 'measured_at_utc': time.strftime('%Y-%m-%dT%H:%M:%SZ', time.gmtime()),
              'environment': {'os': platform.freedesktop_os_release().get('PRETTY_NAME'),
                              'kernel': platform.release(), 'cpu': cpu,
                              'python': platform.python_version(),
                              'zlib': zlib.ZLIB_RUNTIME_VERSION,
                              'xz': subprocess.check_output(['xz', '--version'], text=True).splitlines()[0],
                              'zstd': subprocess.check_output(['zstd', '--version'], text=True).strip()},
              'method': {'repeats': args.repeats, 'warmups_per_configuration': 1,
                         'clock': 'perf_counter_ns around GNU time subprocess; includes process startup and file I/O',
                         'memory': 'GNU time %M, peak RSS in KiB of each compressor/decompressor process',
                         'threads': 'xz -T1; zstd --single-thread',
                         'cache': 'warm cache; no cache flush; decompression output to /dev/null',
                         'order': 'configuration order rotated by one each round',
                         'generator': 'random.Random(20261026), Python randbytes; see benchmark-xz-zstd.py'},
              'datasets': [], 'results': []}
    with tempfile.TemporaryDirectory(prefix='letsgopc-compression-') as tmp:
        base = Path(tmp)
        rng = random.Random(20261026)
        logs = bytearray()
        index = 0
        while len(logs) < 16 * 1024 * 1024:
            logs.extend((f'2026-10-26T12:{index % 60:02d}:00Z level=INFO '
                         f'service=api route=/v1/items/{rng.randrange(10000)} '
                         f'status={rng.choice([200, 200, 200, 404, 500])} '
                         f'latency_ms={rng.randrange(1000)} request_id={rng.getrandbits(64):016x}\n').encode())
            index += 1
        log_bytes = bytes(logs[:16 * 1024 * 1024])
        data = [('synthetic-log', log_bytes),
                ('pseudorandom', rng.randbytes(8 * 1024 * 1024)),
                ('gzip-log', gzip.compress(log_bytes, compresslevel=6, mtime=0))]
        jobs = []
        for name, content in data:
            source = base / name
            source.write_bytes(content)
            digest = hashlib.sha256(content).hexdigest()
            report['datasets'].append({'name': name, 'bytes': len(content), 'sha256': digest})
            for label, command, tool in configs:
                jobs.append({'dataset': name, 'configuration': label, 'command': command,
                             'tool': tool, 'source': source, 'output': base / f'{name}.{label}',
                             'sha256': digest, 'samples': []})

        def measure(command, output):
            rss_file = base / 'rss.txt'
            with output.open('wb') as stream:
                start = time.perf_counter_ns()
                subprocess.run(['/usr/bin/time', '-f', '%M', '-o', str(rss_file), *command],
                               stdout=stream, stderr=subprocess.PIPE, check=True, env=env)
                seconds = (time.perf_counter_ns() - start) / 1e9
            return {'seconds': seconds, 'peak_rss_kib': int(rss_file.read_text().strip())}

        for round_number in range(args.repeats + 1):
            offset = round_number % len(configs)
            ordered = jobs[offset:] + jobs[:offset]
            for job in ordered:
                compressed = measure([*job['command'], str(job['source'])], job['output'])
                decompressed = measure([job['tool'], '-d', '-c', str(job['output'])], Path('/dev/null'))
                if round_number:
                    job['samples'].append({'compress': compressed, 'decompress': decompressed})
            print(f'Completed {"warmup" if not round_number else "round " + str(round_number)}', flush=True)
        for job in jobs:
            restored = subprocess.check_output([job['tool'], '-d', '-c', str(job['output'])], env=env)
            assert hashlib.sha256(restored).hexdigest() == job['sha256']
            entry = {key: job[key] for key in ('dataset', 'configuration', 'command', 'samples')}
            entry['compressed_bytes'] = job['output'].stat().st_size
            entry['verified_sha256'] = job['sha256']
            entry['median'] = {phase: {metric: statistics.median(s[phase][metric] for s in job['samples'])
                                      for metric in ('seconds', 'peak_rss_kib')}
                               for phase in ('compress', 'decompress')}
            report['results'].append(entry)
    args.output.parent.mkdir(parents=True, exist_ok=True)
    args.output.write_text(json.dumps(report, ensure_ascii=False, indent=2) + '\n')
    print(f'Saved {len(report["results"])} verified results to {args.output}')


if __name__ == '__main__':
    main()