#!/usr/bin/python3
"""Host-specific PVE hook: re-enumerate the idle GTX 965M before guest start.

Install on starhe as local:snippets/vdi-shared-gpu-reset.py for VMs 100-102.
An incomplete preparation stays reserved until an administrator inspects it.
"""
import json
import os
from pathlib import Path
import socket
import subprocess
import sys
import time

MEMBERS = (100, 101, 102)
PCI = Path('/sys/bus/pci/devices')
GPU = PCI / '0000:01:00.0'
BRIDGE = PCI / '0000:00:01.0'
STATE_DIR = Path('/run/vdi-shared-gpu')
BOOT_ID = Path('/proc/sys/kernel/random/boot_id')


def stopped():
    for vmid in MEMBERS:
        status = subprocess.check_output(
            ['qm', 'status', str(vmid)], text=True, timeout=10).strip()
        if status != 'status: stopped':
            raise RuntimeError(f'VM {vmid} is not stopped; refusing GPU reset')


def unused_gpu():
    parent = BRIDGE.resolve(strict=True)
    children = sorted(p.name for p in parent.glob('0000:*')
                      if (p / 'vendor').is_file())
    if children != [GPU.name]:
        raise RuntimeError(f'unexpected PCI bridge devices: {children}')
    if ((GPU / 'vendor').read_text().strip() != '0x10de'
            or (GPU / 'device').read_text().strip() != '0x13d9'
            or (GPU / 'driver').resolve(strict=True).name != 'vfio-pci'):
        raise RuntimeError('unexpected GPU identity or driver')
    group = (GPU / 'iommu_group').resolve(strict=True)
    devices = sorted(p.name for p in (group / 'devices').iterdir())
    if devices != [BRIDGE.name, GPU.name]:
        raise RuntimeError(f'unexpected IOMMU group: {devices}')
    result = subprocess.run(['fuser', '/dev/vfio/' + group.name],
                            capture_output=True, timeout=10)
    if result.returncode != 1:
        raise RuntimeError('GPU group is in use or usage could not be checked')


def reenumerate():
    (GPU / 'power/control').write_text('on')
    (GPU / 'remove').write_text('1')
    time.sleep(2)
    (BRIDGE / 'rescan').write_text('1')
    for _ in range(20):
        if GPU.exists() and (GPU / 'driver').exists():
            break
        time.sleep(.5)
    unused_gpu()
    (GPU / 'power/control').write_text('on')
    if (GPU / 'power_state').read_text().strip() != 'D0':
        raise RuntimeError('GPU did not return to D0')


def save(state):
    temporary = STATE_DIR / 'reservation.tmp'
    temporary.write_text(json.dumps(state) + '\n')
    temporary.replace(STATE_DIR / 'reservation.json')


def handle(vmid, phase):
    path = STATE_DIR / 'reservation.json'
    boot = BOOT_ID.read_text().strip()
    state = json.loads(path.read_text()) if path.exists() else None
    if state and state.get('boot') != boot:
        state = None
    if phase == 'pre-start':
        stopped()
        unused_gpu()
        if state and state.get('phase') != 'running':
            raise RuntimeError('previous GPU preparation is unconfirmed; inspect '
                               '/run/vdi-shared-gpu/reservation.json before retry')
        state = dict(vmid=vmid, phase='preparing', boot=boot)
        save(state)
        try:
            reenumerate()
        except Exception:
            save(dict(state, phase='failed'))
            raise
        save(dict(state, phase='prepared'))
        print(f'VDI GPU: exclusive re-enumeration complete for VM {vmid}', flush=True)
    elif phase == 'post-start':
        if not state or state.get('vmid') != vmid or state.get('phase') != 'prepared':
            raise RuntimeError('GPU start reservation does not match this VM')
        save(dict(state, phase='running'))
    elif phase == 'post-stop':
        if state and state.get('vmid') == vmid and state.get('phase') in ('prepared', 'running'):
            stopped()
            unused_gpu()
            path.unlink()
    elif phase != 'pre-stop':
        raise RuntimeError('unsupported hook phase')


def main():
    import fcntl

    if os.geteuid() != 0 or socket.gethostname().split('.')[0] != 'starhe':
        raise RuntimeError('this hook is restricted to root on starhe')
    if len(sys.argv) != 3 or int(sys.argv[1]) not in MEMBERS:
        raise RuntimeError('unexpected VM or hook arguments')
    STATE_DIR.mkdir(mode=0o700, exist_ok=True)
    with (STATE_DIR / 'lock').open('a') as lock:
        fcntl.flock(lock, fcntl.LOCK_EX)
        handle(int(sys.argv[1]), sys.argv[2])


if __name__ == '__main__':
    try:
        main()
    except Exception as error:
        print(f'VDI GPU preparation blocked: {error}', file=sys.stderr, flush=True)
        sys.exit(1)