import apt_pkg
import hashlib
import json
import os
from pathlib import Path
import re
import subprocess

ROOT = Path(__file__).resolve().parents[1]
os.umask(0o077)
checks = {}


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


apt_pkg.init_config()
patterns = apt_pkg.config.value_list('Unattended-Upgrade::Package-Blacklist')
blocked = ['nvidia-driver-595-open', 'nvidia-dkms-595-open', 'nvidia-utils-595',
           'libnvidia-compute-595', 'libnvidia-compute-595:i386', 'libnvidia-container1',
           'nvidia-container-toolkit', 'nvidia-container-toolkit-base',
           'xserver-xorg-video-nvidia-595', 'cuda-drivers', 'cuda-drivers-595',
           'linux-generic', 'linux-image-generic', 'linux-image-6.8.0-146-generic',
           'linux-headers-6.8.0-146-generic', 'linux-modules-extra-6.8.0-146-generic', 'linux-tools-common']
allowed = ['openssl', 'libssl3t64', 'sudo', 'openssh-server', 'curl', 'python3',
           'libc6', 'linux-libc-dev', 'docker-ce']
match = lambda s: any(re.match(p, s) for p in patterns)
assert all(match(n) for n in blocked) and not any(match(n) for n in allowed)
assert apt_pkg.config.find('APT::Periodic::Unattended-Upgrade') == '1'
assert apt_pkg.config.find('APT::Periodic::Update-Package-Lists') == '1'
assert apt_pkg.config.find('Unattended-Upgrade::Automatic-Reboot') == 'false'
checks['apt'] = {'blocked_examples': len(blocked), 'unrelated_examples_allowed': len(allowed),
                 'security_updates': True, 'manual_apt_holds': run('apt-mark', 'showhold')['stdout'],
                 'patterns': patterns}
checks['dpkg_audit'] = run('dpkg', '--audit')
assert checks['dpkg_audit']['rc'] == 0 and not checks['dpkg_audit']['stdout']
checks['systemd_verify'] = run('systemd-analyze', 'verify', 'nvidia-cdi-refresh.service', 'docker.service')
assert checks['systemd_verify']['rc'] == 0
assert all('Ignoring unknown escape sequences' in s for s in checks['systemd_verify']['stderr'].splitlines())
checks['effective_cdi'] = run('systemctl', 'show', 'nvidia-cdi-refresh.service',
                              '--property=ExecCondition', '--property=ExecStartPre', '--property=TimeoutStartUSec')
assert '/usr/sbin/modinfo' in checks['effective_cdi']['stdout']
assert '/usr/bin/grep' not in checks['effective_cdi']['stdout']
assert 'TimeoutStartUSec=45s' in checks['effective_cdi']['stdout']
checks['docker_order'] = run('systemctl', 'show', 'docker.service', '--property=Wants', '--property=After', '--property=Requires')
lines = dict(s.split('=', 1) for s in checks['docker_order']['stdout'].splitlines())
assert 'nvidia-cdi-refresh.service' in lines['Wants'].split()
assert 'nvidia-cdi-refresh.service' in lines['After'].split()
assert 'nvidia-cdi-refresh.service' not in lines['Requires'].split()
checks['grub_syntax'] = run('grub-script-check', '/boot/grub/grub.cfg')
assert checks['grub_syntax']['rc'] == 0
grub = Path('/boot/grub/grub.cfg').read_text()
default_kernel = re.search(r'^\s*linux\s+/boot/vmlinuz-(\S+)', grub, re.M).group(1)
assert default_kernel == '6.8.0-146-generic'
checks['default_kernel'] = default_kernel
checks['grub_environment'] = run('grub-editenv', '/boot/grub/grubenv', 'list')
assert not re.search(r'^next_entry=.+', checks['grub_environment']['stdout'], re.M)
checks['secure_boot'] = run('mokutil', '--sb-state')
assert 'SecureBoot disabled' in checks['secure_boot']['stdout']
checks['module_versions'] = {}
for kernel in ('6.8.0-117-generic', default_kernel):
    version = run('modinfo', '-k', kernel, '-F', 'version', 'nvidia')
    assert version['rc'] == 0 and version['stdout'] == '595.99.02'
    for module in ('nvidia', 'nvidia_uvm', 'nvidia_modeset', 'nvidia_drm'):
        result = run('modprobe', '--show-depends', '--set-version', kernel, module)
        assert result['rc'] == 0, result
    checks['module_versions'][kernel] = version['stdout']
firmware = run('modinfo', '-k', default_kernel, '-F', 'firmware', 'nvidia')['stdout'].splitlines()
assert all((Path('/lib/firmware') / name).is_file() for name in firmware)
checks['firmware_files'] = firmware
checks['initramfs'] = run('lsinitramfs', f'/boot/initrd.img-{default_kernel}')
assert checks['initramfs']['rc'] == 0
checks['initramfs'] = {'readable': True,
    'nvidia_modules_in_initramfs': any('updates/dkms/nvidia' in s for s in checks['initramfs']['stdout'].splitlines()),
    'module_load_path': 'root filesystem DKMS modules via udev/modprobe; no initramfs rebuild performed'}
checks['data_mount'] = run('findmnt', '--fstab', '--evaluate', '--target', '/data', '-o', 'TARGET,SOURCE,FSTYPE,OPTIONS')
assert checks['data_mount']['rc'] == 0 and 'nofail' not in checks['data_mount']['stdout']
checks['ssh_socket_enabled'] = run('systemctl', 'is-enabled', 'ssh.socket')
assert checks['ssh_socket_enabled']['stdout'] == 'enabled'
checks['console_access'] = 'User confirmed Proxmox/physical console access.'
checks['current_loaded_driver'] = Path('/proc/driver/nvidia/version').read_text().splitlines()[0]
checks['runtime_after_reboot'] = 'NOT_RUN; no service restart or GPU module reload performed'
(ROOT / 'reports/host-validation.json').write_text(json.dumps(checks, indent=2) + '\n')
print(json.dumps({'status': 'PREBOOT_CHECKS_PASS', 'default_kernel': default_kernel,
                  'disk_driver': '595.99.02', 'boot_validation': 'NOT_RUN',
                  'systemd_warning': 'vendor base-unit escape warning; overridden ExecCondition verified'}))