#!/usr/bin/env python3
"""Measure z/2/3 stage0 builds on the same full ARM VM; never claim Pi timings."""
import json
import os
from pathlib import Path
import statistics
import subprocess
import sys
import tempfile

project = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(project/'tools'))
from image_formats import digest
from vm_test import VM

if len(sys.argv) != 1: sys.exit('Usage: tools/compare-rust-profiles')
work = Path(tempfile.mkdtemp(prefix='rust-profiles.', dir=project/'out'))
image = (project/'out/fds-system-cli.img').resolve(strict=True)
records = {}
for level in ('z', '2', '3'):
    build = work/('build-'+level)
    environment = dict(os.environ, CARGO_TARGET_DIR=str(build), CARGO_PROFILE_RELEASE_OPT_LEVEL=level)
    with (work/('build-'+level+'.log')).open('wb') as log:
        subprocess.run([str(project/'tools/cargo-build'), '--locked', '--offline', '--release', '--target',
                        'aarch64-unknown-linux-musl', '-p', 'fds-stage0'], cwd=project,
                       env=environment, check=True, stdout=log, stderr=subprocess.STDOUT)
    binary = build/'aarch64-unknown-linux-musl/release/fds-stage0'
    ramfs = work/('initramfs-'+level)
    ramfs.mkdir()
    with (work/('initramfs-'+level+'.log')).open('wb') as log:
        subprocess.run([str(project/'image/build-initramfs'), '--stage0', str(binary),
                        '--output-directory', str(ramfs)], check=True, stdout=log, stderr=subprocess.STDOUT)
    samples = []
    for trial in range(3):
        name = f'opt-{level}-run-{trial+1}'
        with VM(work, name, image, initramfs=ramfs/'initramfs.cpio') as vm:
            vm.expect(rb'FDS> ')
            vm.send('printf "\\nFDS_TRACE_BEGIN\\n"; fds --json boot-profile; printf "\\nFDS_TRACE_END\\n"')
            capture = vm.expect(rb'^FDS_TRACE_BEGIN\r?\n(\{.*?\})\r?\n\r?\nFDS_TRACE_END\r?$')
            report = json.loads(capture.group(1))
            assert report['missing_events'] == ['desktop-ready']
            assert 'virt' in report['platform'].lower()
            (work/(name+'.json')).write_text(json.dumps(report, indent=2)+'\n')
            samples.append(report['durations_ns'])
    records[level] = {'binary_bytes': binary.stat().st_size, 'binary_sha256': digest(binary),
                      'samples_ns': samples,
                      'median_ns': {key: statistics.median(sample[key] for sample in samples) for key in samples[0]}}
    print(f'PASS: opt-level={level}, {binary.stat().st_size} bytes, three measured VM boots', flush=True)
result = {'scope': 'QEMU TCG software comparison; not physical Pi performance', 'samples_per_level': 3,
          'kernel_sha256': digest(project/'out/kernel/boot/kernel_2712.img'),
          'system_sha256': digest(image), 'levels': records,
          'decision': 'Retain opt-level=z until physical Pi measurements justify changing it.'}
(work/'comparison.json').write_text(json.dumps(result, indent=2)+'\n')
target = project/'out/rust-profiles-latest'
temporary = target.with_suffix('.next')
temporary.symlink_to(work.name)
temporary.replace(target)
print(f'PASS: stage0 optimization comparison: {work}')