"""Resume and verify installed-package rollback archives without installing anything."""
import hashlib
import json
import os
from pathlib import Path
import re
import socket
import subprocess
import tarfile

os.umask(0o077)
assert os.geteuid() == 0 and socket.gethostname() == 'icertirnd-ai'
base = Path('/storage/contextra/deploy/ai2-20261006-r01/driver-restore-r01')
packages = base / 'packages'
env = dict(os.environ, LC_ALL='C', DPKG_DEB_THREADS_MAX='2')
before = json.loads((base / 'packages-before.json').read_text())
assert len(before) == 23 and (base / 'configuration-before.tar.gz').is_file()

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

def digest(path):
    with path.open('rb') as f:
        return hashlib.file_digest(f, 'sha256').hexdigest()

def verify_payload(archive, package):
    expected = set()
    for name in command(['dpkg-query', '-L', package]).splitlines():
        if name.startswith('/'):
            p = Path(name)
            if not p.is_dir():
                assert p.is_file() or p.is_symlink(), name
                expected.add(os.path.normpath(name))
        else:
            assert re.fullmatch(r'diverted by base-files to: /(?:lib|lib32|lib64|bin|sbin)\.usr-is-merged', name), name
    actual = set()
    contents = {}
    hardlinks = {}
    with (base / 'archive-reader.log').open('a') as log:
        process = subprocess.Popen(['dpkg-deb', '--fsys-tarfile', str(archive)], stdout=subprocess.PIPE, stderr=log, env=env)
        try:
            with tarfile.open(fileobj=process.stdout, mode='r|') as tar:
                for entry in tar:
                    path = os.path.normpath('/' + entry.name.removeprefix('./').lstrip('/'))
                    assert '..' not in Path(entry.name).parts
                    if entry.isdir():
                        continue
                    actual.add(path)
                    target = Path(path)
                    if entry.issym():
                        assert target.is_symlink() and os.readlink(target) == entry.linkname, path
                    elif entry.islnk():
                        hardlinks[path] = os.path.normpath('/' + entry.linkname.removeprefix('./').lstrip('/'))
                    else:
                        assert entry.isfile(), path
                        data = tar.extractfile(entry)
                        h = hashlib.sha256()
                        while chunk := data.read(1024 * 1024):
                            h.update(chunk)
                        contents[path] = h.hexdigest()
                        assert digest(target) == contents[path], path
            assert process.wait(timeout=30) == 0
        finally:
            process.stdout.close()
            if process.poll() is None:
                process.terminate()
                process.wait(timeout=10)
    for path, link in hardlinks.items():
        visited = set()
        while link in hardlinks:
            assert link not in visited
            visited.add(link)
            link = hardlinks[link]
        assert link in contents and digest(Path(path)) == contents[link], path
    # /lib is a host directory symlink but is represented as a directory in usrmerge packages.
    actual = {p for p in actual if not Path(p).is_dir()}
    assert actual == expected, {'missing': sorted(expected - actual), 'extra': sorted(actual - expected)}
    return len(expected)

results = {}
repack = base / 'backup-tool/dpkg-repack-ai2'
assert digest(repack) == json.loads((base / 'backup-tool/patched-tool.json').read_text())['sha256']
for name, info in before.items():
    assert command(['dpkg-query', '-W', '-f=${Version}', name]).strip() == info['version']
    archive = packages / (name.split(':')[0] + '_' + info['version'] + '_' + info['arch'] + '.deb')
    if not archive.exists():
        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)
        assert r.returncode == 0, name
    log = (base / ('repack-' + name.replace(':', '_') + '.log')).read_text()
    for line in log.splitlines():
        if re.search(r'warning:|error:', line, re.I):
            known_diversion = re.fullmatch(r"dpkg-repack(?:-ai2)?: warning: cannot find file 'diverted by base-files to: /(?:lib|lib32|lib64|bin|sbin)\.usr-is-merged'", line)
            known_field = re.fullmatch(r"dpkg-repack(?:-ai2)?: warning: unknown information field '(Modaliases|Pmaliases|Prefer-Variant|Support)' in input data in stanza in dpkg's status file", line)
            assert known_diversion or known_field, line
    metadata = command(['dpkg-deb', '-f', str(archive), 'Package', 'Version', 'Architecture'])
    assert 'Version: ' + info['version'] in metadata and 'Architecture: ' + info['arch'] in metadata
    for field in ('Modaliases', 'Pmaliases', 'Prefer-Variant', 'Support'):
        assert command(['dpkg-query', '-W', '-f=${' + field + '}', name]).strip() == command(['dpkg-deb', '-f', str(archive), field]).strip(), field
    count = verify_payload(archive, name)
    results[name] = dict(info, archive=archive.name, sha256=digest(archive), size=archive.stat().st_size, verified_payload_files=count)
    (base / 'backup-progress.json').write_text(json.dumps(results, indent=2))
    print('BACKUP_VERIFIED ' + name + ' files=' + str(count), flush=True)
for filename, value in [('rollback-packages.json', results), ('backup-complete.json', {
        'status': 'BACKUP_PACKAGES_READY', 'count': len(results), 'payloads_verified': True,
        'driver_changed': False, 'reboot': False, 'rollback_executed': False,
        'note': 'Pre-change package reconstruction; not a claim of GPU working on kernel139.'})]:
    with (base / filename).open('x') as f:
        json.dump(value, f, indent=2)
print('BACKUP_COMPLETE_NO_DRIVER_CHANGE', flush=True)