# inspired from https://nixos.wiki/wiki/Virtualization_in_NixOS
{ config, pkgs, lib, pkgs-no-overlay, ... }@args:
let
  toImage = f: "${import ./vms/base_image.nix f (args // { myEnv = config.myEnv; })}/nixos.qcow2";
in
{
  options = {
    myServices.vms.libvirt-guests = lib.mkOption {
      type = lib.types.attrsOf (lib.types.submodule {
        options = {
          network = lib.mkOption { type = lib.types.str; description = "network to attach the guest to"; };
          pool = lib.mkOption { type = lib.types.str; description = "pool to attach the guest to"; };
          cpus = lib.mkOption { type = lib.types.int; default = 1; description = "number of cpus to assign"; };
          memory = lib.mkOption { type = lib.types.int; description = "memory in GiB to assign"; };
          diskSize = lib.mkOption { type = lib.types.int; description = "disk size in GiB"; };
          destroyVolumeOnExit = lib.mkOption { type = lib.types.bool; description = "Whether to destroy the volume on exit"; default = false; };
          extraDevicesXML = lib.mkOption { type = lib.types.lines; description = "Extra device configuration"; default = ""; };
          preStart = lib.mkOption { type = lib.types.lines; default = ""; description = "Script to run as prestart"; };
        };
      });
      default = {};
      description = "Libvirt guests to start";
    };
    myServices.vms.libvirt-networks = lib.mkOption {
      type = lib.types.attrsOf (lib.types.submodule {
        options = {
          bridgeNumber = lib.mkOption { type = lib.types.int; description = "bridge interface to create virbr<nbr>"; };
          ipRange = lib.mkOption { type = lib.types.str; example = "192.168.100"; description = "ip4 prefix to use"; };
        };
      });
      description = "Libvirt networks to configure";
      default = {};
    };
    myServices.vms.libvirt-pools = lib.mkOption {
      type = lib.types.attrsOf (lib.types.submodule {
        options = {
          type = lib.mkOption { type = lib.types.enum [ "dir" "zfs" ]; description = "Pool type"; };
          target = lib.mkOption { type = lib.types.nullOr lib.types.path; default = null; description = "where to find images"; };
          preStart = lib.mkOption { type = lib.types.lines; default = ""; description = "Script to run as prestart"; };
          xml = lib.mkOption { type = lib.types.lines; default = ""; description = "Additional configuration"; };
        };
      });
    };
    myServices.vms.libvirt-images = lib.mkOption {
      type = lib.types.attrsOf lib.types.path;
      default = {};
      description = "Attrs of images to create in /etc/libvirtd/base-images";
    };
  };
  config = lib.mkMerge [
    # Define images
    {
      environment.etc = lib.mapAttrs'
        (n: v: lib.nameValuePair "libvirtd/base-images/${n}.qcow2" { source = toImage v; })
        config.myServices.vms.libvirt-images;
    }

    # Define networks
    {
      systemd.services = lib.mapAttrs' (name: network: lib.nameValuePair "libvirtd-network-${name}" {
        after = [ "libvirtd.service" ];
        requires = [ "libvirtd.service" ];
        wantedBy = [ "multi-user.target" ];
        serviceConfig = {
          Type = "oneshot";
          RemainAfterExit = "yes";
        };
        path = [ config.boot.zfs.package ];
        script = let
          xml = pkgs.writeText "libvirt-network-${name}.xml" ''
            <network>
              <name>${name}</name>
              <uuid>UUID</uuid>
              <forward mode='nat' />
              <bridge name='virbr${builtins.toString network.bridgeNumber}' />
              <domain name='${name}' localOnly='yes'/>
              <ip address='${network.ipRange}.1' netmask='255.255.255.0'>
                <dhcp>
                  <range start='${network.ipRange}.2' end='${network.ipRange}.254'/>
                </dhcp>
              </ip>
            </network>
          '';
        in ''
          uuid="$(${pkgs.libvirt}/bin/virsh net-uuid '${name}' || true)"
          ${pkgs.libvirt}/bin/virsh net-define <(sed "s/UUID/$uuid/" '${xml}')
          ${pkgs.libvirt}/bin/virsh net-start '${name}'
        '';
        preStop = ''
          ${pkgs.libvirt}/bin/virsh net-destroy '${name}'
        '';
      }) config.myServices.vms.libvirt-networks;
    }

    # Define pools
    {
      systemd.services = lib.mapAttrs' (name: pool: lib.nameValuePair "libvirtd-pool-${name}" {
        after = [ "libvirtd.service" ];
        requires = [ "libvirtd.service" ];
        wantedBy = [ "multi-user.target" ];
        serviceConfig = {
          Type = "oneshot";
          RemainAfterExit = "yes";
        };
        path = [ config.boot.zfs.package ];
        script = let
          xml = pkgs.writeText "libvirt-pool-${name}.xml" ''
            <pool type="${pool.type}">
              <name>${name}</name>
              <uuid>UUID</uuid>
              ${pool.xml}
              ${if pool.target != null then ''
                <target>
                  <path>${pool.target}</path>
                </target>
              '' else ""}
            </pool>
          '';
        in pool.preStart + ''
          uuid="$(${pkgs.libvirt}/bin/virsh pool-uuid '${name}' || true)"
          ${pkgs.libvirt}/bin/virsh pool-define <(sed "s/UUID/$uuid/" '${xml}')
          ${pkgs.libvirt}/bin/virsh pool-start '${name}' || true
        '';
      }) config.myServices.vms.libvirt-pools;
    }

    # Define guests
    {
      systemd.services = lib.mapAttrs' (name: guest: lib.nameValuePair "libvirtd-guest-${name}" {
        after = [ "libvirtd.service" "libvirtd-pool-${guest.pool}.service" "libvirtd-network-${guest.network}.service" ];
        requires = [ "libvirtd.service" "libvirtd-pool-${guest.pool}.service" "libvirtd-network-${guest.network}.service" ];
        wantedBy = [ "multi-user.target" ];
        serviceConfig = {
          Type = "oneshot";
          RemainAfterExit = "yes";
        };
        path = [ config.boot.zfs.package ];
        script =
          let
            xml = pkgs.writeText "libvirt-guest-${name}.xml"
            ''
              <domain type="kvm">
                <name>${name}</name>
                <uuid>UUID</uuid>
                <memory unit="GiB">${builtins.toString guest.memory}</memory>
                <vcpu>${builtins.toString guest.cpus}</vcpu>
                <os>
                  <type arch="x86_64">hvm</type>
                </os>
                <devices>
                  <emulator>/run/current-system/sw/bin/qemu-system-x86_64</emulator>
                  <disk type="volume">
                    <source pool="${guest.pool}" volume="guest-${name}" />
                    <target dev="vda" bus="virtio"/>
                  </disk>
                  ${guest.extraDevicesXML}
                  <input type="keyboard" bus="usb"/>
                  <graphics type="vnc" port="-1" autoport="yes"/>
                  <interface type="network">
                    <source network="${guest.network}" />
                  </interface>
                </devices>
                <features>
                  <acpi/>
                </features>
              </domain>
            '';
          in
          guest.preStart + ''
            if ! ${pkgs.libvirt}/bin/virsh vol-key 'guest-${name}' --pool ${guest.pool} &> /dev/null; then
              ${pkgs.libvirt}/bin/virsh vol-create-as --pool ${guest.pool} --name 'guest-${name}' --capacity '${builtins.toString guest.diskSize}GiB'
              volume_path=$(${pkgs.libvirt}/bin/virsh vol-path --pool ${guest.pool} --vol 'guest-${name}')
              ${pkgs-no-overlay.qemu}/bin/qemu-img convert /etc/libvirtd/base-images/nixos.qcow2 $volume_path
            fi
            uuid="$(${pkgs.libvirt}/bin/virsh domuuid '${name}' || true)"
            ${pkgs.libvirt}/bin/virsh define <(sed "s/UUID/$uuid/" '${xml}')
            ${pkgs.libvirt}/bin/virsh start '${name}'
          '';
        preStop = ''
          ${pkgs.libvirt}/bin/virsh shutdown '${name}'
          let "timeout = $(date +%s) + 10"
          while [ "$(${pkgs.libvirt}/bin/virsh list --name | grep --count '^${name}$')" -gt 0 ]; do
            if [ "$(date +%s)" -ge "$timeout" ]; then
              # Meh, we warned it...
              ${pkgs.libvirt}/bin/virsh destroy '${name}'
            else
              # The machine is still running, let's give it some time to shut down
              sleep 0.5
            fi
          done
        '' + lib.optionalString guest.destroyVolumeOnExit ''
            if ${pkgs.libvirt}/bin/virsh vol-key 'guest-${name}' --pool ${guest.pool} &> /dev/null; then
              ${pkgs.libvirt}/bin/virsh vol-wipe --pool ${guest.pool} --vol 'guest-${name}' || true
              ${pkgs.libvirt}/bin/virsh vol-delete --pool ${guest.pool} --vol 'guest-${name}'
            fi
        '';
      }) config.myServices.vms.libvirt-guests;
    }
  ];
}