#!/usr/bin/env python3
"""Disposable native benchmark. Does not measure models, OAuth, GUI or viewer."""
import argparse
import datetime
import hashlib
import json
from pathlib import Path
import platform
import subprocess
import time
import uuid

ROOT = Path(__file__).resolve().parents[1]
INNER = r'''
import base64, concurrent.futures, hashlib, json, math, pathlib, statistics, subprocess, tempfile, time
with tempfile.TemporaryDirectory() as directory:
    root = pathlib.Path(directory)
    helper = root / 'computer_files.py'
    helper.write_text(HELPER)
    payload = bytes(range(256)) * 4096
    digest = hashlib.sha256(payload).hexdigest()
    def sample(agent, index):
        started = time.perf_counter()
        name = f'bots/{agent}/sample-{index}' if mode == 'shared' else f'sample-{index}'
        write = subprocess.run(['python3', str(helper), 'write', str(root), name, 'null'],
            input=base64.b64encode(payload), capture_output=True, timeout=15)
        read = subprocess.run(['python3', str(helper), 'read', str(root), name], capture_output=True, timeout=15)
        command = subprocess.run(['python3', '-c', 'print(sum(range(100000)))'], capture_output=True, timeout=15)
        ok = write.returncode == read.returncode == command.returncode == 0
        ok = ok and hashlib.sha256(read.stdout).hexdigest() == digest and command.stdout == b'4999950000\n'
        return {'agent': agent, 'index': index, 'ok': ok, 'elapsed_ms': (time.perf_counter()-started)*1000}
    agents = ['a','b'] if mode == 'shared' else ['c']
    started = time.perf_counter()
    with concurrent.futures.ThreadPoolExecutor(max_workers=len(agents)) as pool:
        results = list(pool.map(lambda item: sample(*item), [(a,i) for i in range(samples) for a in agents]))
    durations = sorted(r['elapsed_ms'] for r in results)
    cgroup = {}
    for metric in ('memory.peak', 'memory.max', 'cpu.max', 'cpu.stat', 'pids.max'):
        path = pathlib.Path('/sys/fs/cgroup') / metric
        if path.exists():
            cgroup[metric] = path.read_text().strip()
    print(json.dumps({'mode':mode, 'agents':agents, 'samples':results,
        'successes':sum(r['ok'] for r in results), 'count':len(results),
        'median_ms':statistics.median(durations), 'p95_ms':durations[math.ceil(len(durations)*.95)-1],
        'wall_ms':(time.perf_counter()-started)*1000, 'cgroup':cgroup,
        'desktop_started':False, 'viewer_started':False}))
'''


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('--image', default='lazyboy/computer:local')
    parser.add_argument('--samples', type=int, default=30)
    parser.add_argument('--output', default='docs/benchmarks/agent-computer-native.json')
    args = parser.parse_args()
    if not 2 <= args.samples <= 100:
        parser.error('samples must be 2..100')
    image_id = subprocess.check_output(['docker', 'image', 'inspect', '--format', '{{.Id}}', args.image], text=True).strip()
    report = {'timestamp': datetime.datetime.now(datetime.timezone.utc).isoformat(),
              'image': args.image, 'image_id': image_id, 'host': platform.platform(),
              'scope': 'Computer-local helper write/read/hash + native exec; 1 MiB per sample; Team A/B concurrent, Dedicated C alone',
              'limitations': ['No API/harness/model timing', 'No real OAuth', 'No viewer/GUI latency or readability measurement', 'No baseline speedup claim'],
              'results': []}
    helper = (ROOT / 'crates/supervisor/src/computer_files.py').read_text()
    report['helper_sha256'] = hashlib.sha256(helper.encode()).hexdigest()
    for mode in ('shared', 'dedicated'):
        name = 'lazyboy-plan-bench-' + uuid.uuid4().hex[:12]
        started = time.perf_counter()
        try:
            subprocess.run(['docker', 'run', '-d', '--rm', '--name', name, '--network', 'none',
                            '--cpus', '1', '--memory', '512m', '--pids-limit', '128',
                            '--entrypoint', 'sleep', args.image, '300'], check=True, capture_output=True)
            provision_ms = (time.perf_counter() - started) * 1000
            source = f'HELPER={helper!r}\nsamples={args.samples}\nmode={mode!r}\n' + INNER
            result = subprocess.run(['docker', 'exec', '-i', '--user', '1000:1000', name, 'python3', '-'],
                                    input=source, capture_output=True, text=True, timeout=240, check=True)
            row = json.loads(result.stdout)
            row['container_start_ms'] = provision_ms
            report['results'].append(row)
            print(f"{mode}: {row['successes']}/{row['count']} success; median {row['median_ms']:.1f} ms; P95 {row['p95_ms']:.1f} ms", flush=True)
        finally:
            subprocess.run(['docker', 'rm', '-f', name], capture_output=True)
    output = Path(args.output)
    output.parent.mkdir(parents=True, exist_ok=True)
    output.write_text(json.dumps(report, indent=2) + '\n')
    if any(row['successes'] != row['count'] for row in report['results']):
        raise SystemExit(1)


if __name__ == '__main__':
    main()