nixos/modules/winvm.nix

{ den, ... }:
{
  den.aspects.winvm.nixos = { config, pkgs, lib, ... }:
  let
    cfg = config.local.winvm;

    hookScript = pkgs.writeShellScript "qemu-hook" ''
      #!/usr/bin/env bash
      # Libvirt QEMU hook for single-GPU passthrough.
      # Libvirt calls this script as: qemu <vm-name> <operation> <sub-operation>
      # On "prepare begin": unbind GPU from NVIDIA and hand it to the VM.
      # On "release end": return GPU to the host and restart the display session.

      VM_NAME="$1"
      OPERATION="$2"

      GPU_ADDR="${cfg.gpuPcieAddress}"
      GPU_AUDIO_ADDR="${cfg.gpuAudioPcieAddress}"

      # Exit early for any VM that is not the configured passthrough target
      if [[ "$VM_NAME" != "${cfg.vmName}" ]]; then
          exit 0
      fi

      bind_vfio() {
          # Stop greetd before terminating the session so it cannot restart tuigreet
          systemctl stop greetd.service

          # Terminate the Niri/Wayland session and all user services (including Pipewire)
          loginctl terminate-user zoty 2>/dev/null || true
          sleep 3

          # Detach virtual console framebuffer bindings held by the NVIDIA driver
          echo 0 > /sys/class/vtconsole/vtcon0/bind 2>/dev/null || true
          echo 0 > /sys/class/vtconsole/vtcon1/bind 2>/dev/null || true

          # Detach the EFI framebuffer platform driver if it is present
          echo efi-framebuffer.0 \
              > /sys/bus/platform/drivers/efi-framebuffer/unbind 2>/dev/null || true

          sleep 1

          # Unload NVIDIA kernel modules in reverse dependency order
          modprobe -r nvidia_drm nvidia_modeset nvidia_uvm nvidia

          # Bind GPU and its HDMI audio function to the vfio-pci driver
          echo "$GPU_ADDR"       > /sys/bus/pci/drivers/vfio-pci/bind
          echo "$GPU_AUDIO_ADDR" > /sys/bus/pci/drivers/vfio-pci/bind
      }

      unbind_vfio() {
          # Release both GPU PCI functions from vfio-pci
          echo "$GPU_ADDR"       > /sys/bus/pci/drivers/vfio-pci/unbind
          echo "$GPU_AUDIO_ADDR" > /sys/bus/pci/drivers/vfio-pci/unbind

          sleep 1

          # Reload NVIDIA modules in dependency order
          modprobe nvidia
          modprobe nvidia_modeset
          modprobe nvidia_uvm
          modprobe nvidia_drm

          # Restore the EFI framebuffer driver
          echo efi-framebuffer.0 \
              > /sys/bus/platform/drivers/efi-framebuffer/bind 2>/dev/null || true

          # Restore virtual console framebuffer bindings
          echo 1 > /sys/class/vtconsole/vtcon0/bind 2>/dev/null || true
          echo 1 > /sys/class/vtconsole/vtcon1/bind 2>/dev/null || true

          # Restart greetd so tuigreet is available for login
          systemctl start greetd.service
      }

      case "$OPERATION" in
          prepare) bind_vfio   ;;
          release) unbind_vfio ;;
      esac
    '';

    screamService = {
      Unit = {
        Description = "Scream audio receiver for Windows VM";
        # Start after Pipewire is ready so the sink exists when scream connects
        After = [ "pipewire-pulse.service" ];
        Requires = [ "pipewire-pulse.service" ];
      };
      Service = {
        # virbr0 is the default libvirt NAT bridge; scream listens for UDP audio
        # streams from Windows and forwards them to the PulseAudio/Pipewire sink
        ExecStart = "${pkgs.scream}/bin/scream -i virbr0 -o pulse";
        Restart = "on-failure";
      };
      Install.WantedBy = [ "default.target" ];
    };
  in
  {
    options.local.winvm = {
      gpuPcieAddress = lib.mkOption {
        type        = lib.types.str;
        description = "PCIe address of the NVIDIA GPU (e.g. 0000:01:00.0).";
      };
      gpuAudioPcieAddress = lib.mkOption {
        type        = lib.types.str;
        description = "PCIe address of the NVIDIA GPU HDMI audio device (e.g. 0000:01:00.1).";
      };
      vmName = lib.mkOption {
        type        = lib.types.str;
        default     = "windows";
        description = "Libvirt VM name that triggers the GPU handoff hook.";
      };
    };

    config = {
      boot.kernelParams = [ "intel_iommu=on" "iommu=pt" ];

      boot.kernelModules = [ "vfio" "vfio_pci" "vfio_iommu_type1" ];

      virtualisation.libvirtd = {
        enable = true;
        qemu = {
          package      = pkgs.qemu_kvm;
          runAsRoot    = true;
          swtpm.enable = true;
        };
      };

      programs.virt-manager.enable = true;

      environment.systemPackages = [ pkgs.virtio-win ];

      users.users.zoty.extraGroups = [ "libvirtd" "kvm" ];

      environment.etc."libvirt/hooks/qemu" = {
        source = hookScript;
        mode   = "0755";
      };

      # Scream receiver runs as a user service so it has access to the
      # Pipewire session. It starts automatically after login and waits
      # for audio packets from the Windows VM on the libvirt bridge.
      home-manager.users.zoty.systemd.user.services.scream = screamService;
    };
  };
}