"""Prepare rollback packages only; never installs or loads a driver."""
import hashlib
import json
import os
from pathlib import Path
import re
import shutil
import socket
import subprocess
import urllib.request

os.umask(0o077)
assert os.geteuid() == 0 and socket.gethostname() == 'icertirnd-ai'
env = dict(os.environ, LC_ALL='C', DPKG_DEB_THREADS_MAX='2')
base = Path('/storage/contextra/deploy/ai2-20261006-r01/driver-restore-r01')

def command(args, timeout=30):
    r = subprocess.run(args, capture_output=True, text=True, timeout=timeout, env=env)
    if r.returncode:
        raise RuntimeError(str(args[0]) + ' failed: ' + r.stderr[-2000:])
    return r.stdout

def save(name, value):
    path = base / name
    with path.open('x') as f:
        json.dump(value, f, ensure_ascii=False, indent=2)
        f.write('\n')

assert command(['uname', '-r']).strip() == '6.8.0-139-generic'
assert not command(['docker', 'ps', '-q']).strip()
assert shutil.disk_usage('/').free > 15 * 1024**3
assert shutil.disk_usage('/storage').free > 20 * 1024**3
mount = json.loads(command(['findmnt', '-J', '-o', 'UUID', '/storage']))
assert mount['filesystems'][0]['uuid'] == 'e49118da-66c6-49c3-b180-1d9146742b0e'
requested = ['linux-modules-nvidia-595-open-6.8.0-139-generic=6.8.0-139.139+1',
             'nvidia-driver-595-open=595.91.07-0ubuntu0.24.04.1']
raw = command(['apt-get', '-s', '--no-install-recommends', 'install', *requested], 60)
lines = [line for line in raw.splitlines() if line.startswith(('Inst ', 'Remv '))]
inst = [line for line in lines if line.startswith('Inst ')]
removed = {line.split()[1] for line in lines if line.startswith('Remv ')}
assert len(inst) == 23 and sum('[' in line.split('(')[0] for line in inst) == 21
assert removed == {'linux-modules-nvidia-595-open-6.8.0-134-generic', 'linux-modules-nvidia-595-open-6.8.0-136-generic'}
for line in inst:
    package = line.split()[1]
    assert re.fullmatch(r'(?:libnvidia-[a-z0-9-]+-595(?::i386)?|nvidia-[a-z0-9.-]+|linux-modules-nvidia-595-open-(?:generic|6\.8\.0-139-generic)|xserver-xorg-video-nvidia-595)', package)
    version = re.search(r'\((\S+)', line).group(1)
    assert version == ('6.8.0-139.139+1' if package.startswith('linux-modules-') else '595.91.07-0ubuntu0.24.04.1')
plan = {'requests': requested, 'operations': lines}
if base.exists():
    assert not base.is_symlink()
    assert set(p.name for p in base.iterdir()) == {'approved-plan.json', 'apt-simulation.txt'}
    assert json.loads((base / 'approved-plan.json').read_text()) == plan
else:
    base.mkdir(mode=0o700)
    save('approved-plan.json', plan)
    (base / 'apt-simulation.txt').write_text(raw)
before = {}
native = command(['dpkg', '--print-architecture']).strip()
for line in lines:
    if line.startswith('Remv ') or '[' in line.split('(')[0]:
        name = line.split()[1]
        output = command(['dpkg-query', '-W', '-f=${binary:Package}\t${Version}\t${Architecture}\t${Status}\n', name])
        canonical, info = select_installed(output, name, native)
        before[canonical] = info
assert len(before) == 23
save('packages-before.json', before)
save('apt-marks-before.json', {'auto': command(['apt-mark', 'showauto']).splitlines(),
                             'manual': command(['apt-mark', 'showmanual']).splitlines(),
                             'hold': command(['apt-mark', 'showhold']).splitlines()})
config = ['etc/modprobe.d', 'etc/modules-load.d', 'etc/initramfs-tools', 'etc/default/grub',
          'etc/alternatives', 'var/lib/dpkg/status', 'var/lib/dpkg/alternatives',
          'boot/initrd.img-6.8.0-139-generic', 'boot/vmlinuz-6.8.0-139-generic']
config = [p for p in config if Path('/' + p).exists()]
command(['tar', '--acls', '--xattrs', '--numeric-owner', '-czf', str(base / 'configuration-before.tar.gz'), '-C', '/', *config], 180)
tool = base / 'backup-tool'
tool.mkdir(mode=0o700)
deb = tool / 'dpkg-repack_1.52_all.deb'
with urllib.request.urlopen('https://archive.ubuntu.com/ubuntu/pool/main/d/dpkg-repack/dpkg-repack_1.52_all.deb', timeout=20) as response:
    payload = response.read(100000)
assert hashlib.sha256(payload).hexdigest() == '6b36883d88d5b6cf36f59e58734dfaa44e103b2bc44bcbe5f17e82e84a3c5432'
deb.write_bytes(payload)
command(['dpkg-deb', '-x', str(deb), str(tool / 'extracted')])
repack = tool / 'extracted/usr/bin/dpkg-repack'
assert repack.is_file()
packages = base / 'packages'
packages.mkdir(mode=0o700)
results = {}
for name, info in before.items():
    print('BACKUP_START ' + name, flush=True)
    args = ['nice', '-n', '10', 'perl', str(repack), '--tag=none', '--deb-option=-Zzstd', '--deb-option=-z1', name]
    r = subprocess.run(args, cwd=packages, capture_output=True, text=True, timeout=600, env=env)
    (base / ('repack-' + name.replace(':', '_') + '.log')).write_text(r.stdout + r.stderr)
    if r.returncode or re.search(r'warning:|error:', r.stderr, re.I):
        raise RuntimeError('Repack failed or warned for ' + name + ': ' + r.stderr[-1500:])
    expected = packages / (name.split(':')[0] + '_' + info['version'] + '_' + info['arch'] + '.deb')
    assert expected.is_file(), name
    metadata = command(['dpkg-deb', '-f', str(expected), 'Package', 'Version', 'Architecture'])
    assert 'Version: ' + info['version'] in metadata and 'Architecture: ' + info['arch'] in metadata
    with expected.open('rb') as f:
        digest = hashlib.file_digest(f, 'sha256').hexdigest()
    results[name] = dict(info, archive=expected.name, sha256=digest, size=expected.stat().st_size)
    print('BACKUP_VERIFIED ' + name, flush=True)
save('rollback-packages.json', results)
assert len(results) == 23
save('backup-complete.json', {'status': 'BACKUP_PACKAGES_READY', 'count': len(results),
                            'driver_changed': False, 'reboot': False,
                            'rollback_executed': False,
                            'note': 'Repacked installed packages and protected config backup. Original baseline has no GPU module for kernel139.'})
print('BACKUP_COMPLETE_NO_DRIVER_CHANGE', flush=True)