dotfiles/modules/nixos/guests/comfyui/default.nix

{
  flake,
  ...
}:
let
  inherit (flake.config.people) user0;
  inherit (flake.config.services) instances;
  serviceCfg = instances.comfyui;
  hostCfg = instances.web;
in
{
  microvm.vms = {
    ${serviceCfg.name} = {
      autostart = true;
      restartIfChanged = true;
      config = {
        system.stateVersion = "24.05";
        time.timeZone = "America/Winnipeg";
        users.users.root.openssh.authorizedKeys.keys = flake.config.people.users.${user0}.sshKeys;

        imports = [
          flake.inputs.filesorter.nixosModules.default
        ];

        services = {
          filesorter.enable = true;
          caddy = {
            enable = true;
            virtualHosts = {
              ":80" = {
                extraConfig = ''
                  # Main reverse proxy with WebSocket support
                  reverse_proxy ${hostCfg.localhost.address1}:${toString serviceCfg.ports.port0} {
                    header_up Host {host}
                    header_up X-Real-IP {remote}
                    header_up X-Forwarded-For {remote}
                    header_up X-Forwarded-Proto {scheme}
                    
                    # WebSocket support - critical for ComfyUI real-time updates
                    header_up Connection {>Connection}
                    header_up Upgrade {>Upgrade}
                    
                    # Longer timeouts for generation tasks
                    transport http {
                      read_timeout 300s
                      write_timeout 300s
                    }
                  }

                  # Security headers
                  header {
                    Strict-Transport-Security "max-age=31536000; includeSubDomains; preload"
                    X-Frame-Options "DENY"
                    X-Content-Type-Options "nosniff"
                    X-XSS-Protection "1; mode=block"
                    Referrer-Policy "strict-origin-when-cross-origin"
                    -Server
                  }

                  # Logging
                  log {
                    output file /var/log/caddy/comfyui-access.log
                    format json
                  }
                '';
              };
            };
          };
          openssh = {
            enable = true;
            settings = {
              PasswordAuthentication = false;
              PermitRootLogin = "prohibit-password";
            };
          };
        };

        virtualisation = {
          docker = {
            enable = true;
            autoPrune = {
              enable = true;
              dates = "weekly";
            };
          };
          oci-containers = {
            backend = "docker";
            containers.comfyui = {
              image = "yanwk/comfyui-boot:cu126-slim";
              autoStart = true;

              ports = [
                "${hostCfg.localhost.address1}:${builtins.toString serviceCfg.ports.port0}:8188"
              ];

              volumes = [
                "${serviceCfg.varPaths.path0}:/root"
                "${serviceCfg.varPaths.path0}/models:/root/models"
                "${serviceCfg.varPaths.path0}/custom_nodes:/root/custom_nodes"
                "${serviceCfg.varPaths.path0}/output:/root/output"
                "${serviceCfg.varPaths.path0}/input:/root/input"
                "${serviceCfg.varPaths.path0}/user:/root/user"
                # Mount the device nodes into the container
                "/dev/nvidia0:/dev/nvidia0"
                "/dev/nvidiactl:/dev/nvidiactl"
                "/dev/nvidia-uvm:/dev/nvidia-uvm"
                "/dev/nvidia-modeset:/dev/nvidia-modeset"
              ];

              environment = {
                CLI_ARGS = "--listen 0.0.0.0 --port 8188 --preview-method auto --dont-print-server";
                NVIDIA_VISIBLE_DEVICES = "all";
                NVIDIA_DRIVER_CAPABILITIES = "compute,utility,graphics";
              };

              extraOptions = [
                "--privileged"
                "--device=/dev/nvidia0"
                "--device=/dev/nvidiactl"
                "--device=/dev/nvidia-uvm"
                "--device=/dev/nvidia-modeset"
                "--memory=32g"
                "--memory-swap=32g"
                "--shm-size=16g"
                "--security-opt=no-new-privileges:true"
                "--network=bridge"
                "--health-cmd=curl -f http://localhost:8188/ || exit 1"
                "--health-interval=30s"
                "--health-timeout=10s"
                "--health-retries=3"
              ];
            };
          };
        };

        users.users.caddy.extraGroups = [ serviceCfg.name ];

        networking.firewall.allowedTCPPorts = [
          22
          80
        ];

        systemd = {
          services.docker-comfyui = {
            after = [
              "network-online.target"
              "docker.service"
            ];
            wants = [ "network-online.target" ];
            requires = [ "docker.service" ];
          };
          network = {
            enable = true;
            networks."20-lan" = {
              matchConfig.Name = "enp0s5";
              addresses = [ { Address = "${serviceCfg.interface.ip}/24"; } ];
              routes = [
                {
                  Destination = "${hostCfg.localhost.address1}/0";
                  Gateway = serviceCfg.interface.gate;
                }
              ];
              dns = [
                "1.1.1.1"
                "8.8.8.8"
              ];
            };
          };
          tmpfiles.rules = [
            "Z ${serviceCfg.varPaths.path0} 0755 ${serviceCfg.name} ${serviceCfg.name} -"
            "d /var/lib/docker/tmp/ 755 docker docker -"
            "d /var/log/caddy 755 caddy caddy -"
          ];
        };

        boot.kernel.sysctl = {
          "kernel.shmmax" = 68719476736;
          "kernel.shmall" = 4194304;
          "vm.swappiness" = 1;
          "vm.dirty_ratio" = 15;
          "vm.dirty_background_ratio" = 5;
        };

        microvm = {
          vcpu = 6;
          mem = 8192;
          hypervisor = "qemu";
          interfaces = [
            {
              type = "tap";
              id = serviceCfg.interface.id;
              mac = serviceCfg.interface.mac;
            }
            {
              type = "user";
              id = serviceCfg.interface.idUser;
              mac = serviceCfg.interface.macUser;
            }
          ];

          forwardPorts = [
            {
              from = "host";
              host.port = serviceCfg.interface.ssh;
              guest.port = 22;
            }
          ];

          shares = [
            {
              mountPoint = "/nix/.ro-store";
              proto = "virtiofs";
              source = "/nix/store";
              tag = "read_only_nix_store";
            }
            {
              mountPoint = "/var/lib/${serviceCfg.name}";
              proto = "virtiofs";
              source = "${serviceCfg.mntPaths.path0}/data";
              tag = "${serviceCfg.name}_data";
            }
            {
              mountPoint = "/var/lib/docker";
              proto = "virtiofs";
              source = "${serviceCfg.mntPaths.path0}/docker"; # Store on host
              tag = "docker_data";
            }
          ];
        };
      };
    };
  };

  systemd.tmpfiles.rules = [
    "d ${serviceCfg.mntPaths.path0}/data 0751 microvm wheel - -"
    "d ${serviceCfg.mntPaths.path0}/docker 0751 microvm wheel - -"

  ];
}