import hashlib
import json
import os
from pathlib import Path
import shutil
import subprocess
import sys
import time

ROOT = Path(__file__).resolve().parents[1]
os.umask(0o077)
label = sys.argv[1]
assert label in ('before', 'after')


def run(*args):
    result = subprocess.run(args, capture_output=True, text=True, timeout=15)
    return {'rc': result.returncode, 'stdout': result.stdout.strip(), 'stderr': result.stderr.strip()}


files = ['/etc/apt/apt.conf.d/20auto-upgrades', '/etc/apt/apt.conf.d/50unattended-upgrades',
         '/etc/systemd/system/nvidia-cdi-refresh.service', '/etc/systemd/system/nvidia-cdi-refresh.path',
         '/etc/default/grub', '/boot/grub/grubenv', '/etc/nvidia-container-runtime/config.toml',
         '/etc/fstab', '/run/cdi/ssims-emergency-59584.json',
         '/run/systemd/system/docker.service.d/90-contextra-export-temp.conf']
changed = ['/etc/apt/apt.conf.d/52test-gpu-maintenance',
           '/etc/systemd/system/nvidia-cdi-refresh.service.d/20-test-gpu-readiness.conf',
           '/etc/systemd/system/docker.service.d/20-test-gpu-cdi-order.conf']
hashes = {}
for value in files + changed:
    p = Path(value)
    if not p.exists():
        hashes[value] = None
        continue
    assert p.is_file() and not p.is_symlink(), value
    hashes[value] = hashlib.sha256(p.read_bytes()).hexdigest()
    dest = ROOT / 'host-config' / label / value.lstrip('/')
    dest.parent.mkdir(parents=True, exist_ok=True)
    shutil.copy2(p, dest)
    dest.chmod(0o600)
    if label == 'before' and value.startswith('/run/'):
        dest = ROOT / 'files/run-selected' / value.removeprefix('/run/')
        dest.parent.mkdir(parents=True, exist_ok=True)
        shutil.copy2(p, dest)
        dest.chmod(0o600)
if label == 'before':
    assert all(hashes[p] is None for p in changed), 'Existing target: do not overwrite'
ids = run('docker', 'ps', '-aq')['stdout'].split()
containers = json.loads(run('docker', 'inspect', *ids)['stdout']) if ids else []
summary = [{'name': d['Name'].lstrip('/'), 'image_id': d['Image'], 'pid': d['State']['Pid'],
            'running': d['State']['Running'], 'started_at': d['State']['StartedAt'],
            'restart_policy': d['HostConfig']['RestartPolicy'],
            'project': d['Config'].get('Labels', {}).get('com.docker.compose.project')}
           for d in containers]
data = {'captured_epoch': time.time(), 'files': hashes, 'containers': summary,
        'boot_id': Path('/proc/sys/kernel/random/boot_id').read_text().strip(),
        'kernel': os.uname().release,
        'loaded_nvidia': Path('/proc/driver/nvidia/version').read_text().splitlines()[0],
        'apt_policy': run('apt-config', 'dump'),
        'units': run('systemctl', 'show', 'docker.service', 'nvidia-cdi-refresh.service',
                    '--property=After', '--property=Wants', '--property=Requires',
                    '--property=ExecStartPre', '--property=ExecCondition', '--property=TimeoutStartUSec'),
        'git': {p: {'head': run('git', '-C', p, 'rev-parse', 'HEAD'),
                    'status': run('git', '-C', p, 'status', '--short', '--branch')}
                for p in ('/workspace/ssims.ai', '/workspace/contextra-analytics-runtime-lifecycle-refactor')}}
out = ROOT / 'reports' / f'host-{label}.json'
assert not out.exists(), 'Snapshot already exists'
out.write_text(json.dumps(data, indent=2) + '\n')
print(json.dumps({'snapshot': str(out), 'containers': len(summary), 'tracked_targets': len(hashes)}))