"""Apply only the reviewed NVIDIA package plan; no reboot or kernel replacement."""
import hashlib
import json
import os
from pathlib import Path
import shutil
import signal
import socket
import subprocess

os.umask(0o077)
base = Path('/storage/contextra/deploy/ai2-20261006-r01/driver-restore-r01')
env = dict(os.environ, LC_ALL='C', DEBIAN_FRONTEND='noninteractive', NEEDRESTART_MODE='l')

def command(args, timeout=30):
    return subprocess.check_output(args, text=True, stderr=subprocess.STDOUT, env=env, timeout=timeout)

def save(name, value):
    with (base / name).open('x') as f:
        json.dump(value, f, indent=2)

def logged(args, name, seconds):
    with (base / (name + '.log')).open('x') as log:
        process = subprocess.Popen(args, stdout=log, stderr=log, env=env, start_new_session=True)
        try:
            code = process.wait(timeout=seconds)
        except subprocess.TimeoutExpired:
            os.killpg(process.pid, signal.SIGINT)
            try:
                code = process.wait(timeout=30)
            except subprocess.TimeoutExpired:
                os.killpg(process.pid, signal.SIGTERM)
                code = process.wait(timeout=15)
            save(name + '-timeout.json', {'timeout': seconds, 'returncode': code})
            raise RuntimeError(name + ' timed out; inspect package state before any retry')
    save(name + '-result.json', {'returncode': code})
    if code:
        raise RuntimeError(name + ' failed; inspect its protected log')

assert os.geteuid() == 0 and socket.gethostname() == 'icertirnd-ai'
assert command(['uname', '-r']).strip() == '6.8.0-139-generic'
assert not command(['docker', 'ps', '-q']).strip()
assert not command(['dpkg', '--audit']).strip()
assert shutil.disk_usage('/').free > 15 * 1024**3
assert json.loads((base / 'backup-complete.json').read_text())['payloads_verified'] is True
backups = json.loads((base / 'rollback-packages.json').read_text())
assert len(backups) == 23
for name, row in backups.items():
    assert command(['dpkg-query', '-W', '-f=${Version}', name]).strip() == row['version']
    with (base / 'packages' / row['archive']).open('rb') as f:
        assert hashlib.file_digest(f, 'sha256').hexdigest() == row['sha256']
plan = json.loads((base / 'approved-plan.json').read_text())
args = ['apt-get', '--no-install-recommends', '-o', 'Dpkg::Options::=--force-confold',
        '-o', 'APT::Get::AutomaticRemove=false', 'install', *plan['requests']]
raw = command(['apt-get', '-s', '--no-install-recommends', 'install', *plan['requests']], 60)
assert [x for x in raw.splitlines() if x.startswith(('Inst ', 'Remv '))] == plan['operations']
logged([*args[:1], '--assume-yes', '--download-only', *args[1:]], 'download-r02', 600)
deb = Path('/var/cache/apt/archives/linux-modules-nvidia-595-open-6.8.0-139-generic_6.8.0-139.139+1_amd64.deb')
assert deb.is_file()
extract = base / 'new-module-inspection'
assert not extract.exists()
command(['dpkg-deb', '-x', str(deb), str(extract)], 60)
modules = list(extract.rglob('nvidia.ko')) + list(extract.rglob('nvidia.ko.zst'))
assert len(modules) == 1
module = modules[0]
properties = {key: command(['modinfo', '-F', key, str(module)]).strip() for key in ('version', 'vermagic', 'signer', 'sig_id')}
assert properties['version'] == '595.91.07'
assert properties['vermagic'].startswith('6.8.0-139-generic ')
assert properties['signer'] and properties['sig_id'] == 'PKCS#7'
save('module-before-install.json', properties)
print('DOWNLOAD_AND_MODULE_INSPECTION_PASS', flush=True)
raw = command(['apt-get', '-s', '--no-install-recommends', 'install', *plan['requests']], 60)
assert [x for x in raw.splitlines() if x.startswith(('Inst ', 'Remv '))] == plan['operations']
save('install-intent.json', {'approved_requests': plan['requests'], 'kernel': '6.8.0-139-generic', 'reboot': False})
logged([*args[:1], '--assume-yes', *args[1:]], 'install', 900)
assert command(['uname', '-r']).strip() == '6.8.0-139-generic'
assert not command(['dpkg', '--audit']).strip()
assert command(['dpkg-query', '-W', '-f=${Version}', 'nvidia-driver-595-open']).strip() == '595.91.07-0ubuntu0.24.04.1'
save('install-complete.json', {'driver': '595.91.07', 'kernel': '6.8.0-139-generic', 'reboot': False})
print('DRIVER_PACKAGES_INSTALLED_NO_REBOOT', flush=True)
command(['modprobe', 'nvidia'], 30)
command(['modprobe', 'nvidia_uvm'], 30)
gpu = command(['nvidia-smi', '--query-gpu=name,uuid,driver_version,memory.total,memory.used', '--format=csv'], 15)
save('gpu-readback.json', {'nvidia_smi': gpu, 'reboot': False})
print(gpu, flush=True)