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

{
  flake,
  config,
  ...
}:
let
  inherit (flake.config.services) instances;
  serviceCfg = instances.comfyui;
  localhost = instances.web.localhost.address1;
  host = serviceCfg.domains.url0;
  dns = instances.web.dns.provider0;
  dnsPath = "dns/${dns}";

in
{
  virtualisation.docker = {
    enable = true;
    enableNvidia = true;
    autoPrune = {
      enable = true;
      dates = "weekly";
    };
  };

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

      ports = [
        "${localhost}:${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"
      ];

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

      extraOptions = [
        "--runtime=nvidia"
        "--gpus=device=0"
        # Memory limits to prevent OOM
        "--memory=32g"
        "--memory-swap=32g"
        "--shm-size=16g"
        # Security
        "--security-opt=no-new-privileges:true"
        # Network
        "--network=bridge"
        # Health check
        "--health-cmd=curl -f http://localhost:8188/ || exit 1"
        "--health-interval=30s"
        "--health-timeout=10s"
        "--health-retries=3"
      ];
    };
  };

  services.caddy = {
    enable = true;
    virtualHosts = {
      "${host}" = {
        extraConfig = ''
          basic_auth {
            {$CADDY_AUTH_USER} {$CADDY_AUTH_PASSWORD_HASH}
          }

          # Main reverse proxy with WebSocket support
          reverse_proxy ${localhost}:${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
            }
          }

          tls ${serviceCfg.ssl.cert} ${serviceCfg.ssl.key}

          # 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
          }
        '';
      };
    };
  };

  systemd.tmpfiles.rules = [
    "d ${serviceCfg.varPaths.path0} 755 root root -"
    "d ${serviceCfg.secretPaths.path0}/caddy 755 caddy caddy -"
    "d /var/log/caddy 755 caddy caddy -"
  ];

  systemd.services.caddy = {
    serviceConfig = {
      EnvironmentFile = config.sops.secrets."caddy/${serviceCfg.name}-auth".path;
    };
  };

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

  systemd.services.docker-comfyui = {
    after = [
      "network-online.target"
      "docker.service"
    ];
    wants = [ "network-online.target" ];
    requires = [ "docker.service" ];
  };

  sops.secrets = {
    "caddy/${serviceCfg.name}-auth" = {
      owner = "caddy";
      group = "caddy";
      mode = "0400";
    };
  };

  users.users.caddy.extraGroups = [ "acme" ];

  security.acme.certs."${host}" = {
    dnsProvider = dns;
    environmentFile = config.sops.secrets.${dnsPath}.path;
    group = "caddy";
  };

}