{ pkgs, config, lib, ... }:

let
  comfyuiSrc = pkgs.fetchFromGitHub {
    owner = "Comfy-Org";
    repo = "ComfyUI";
    rev = "master";
    sha256 = "PQfZ0PD/PQn49ElGdzt/El8JrU7clETcLi/6ZUxm8f8=";
  };

  dataDir = "${config.home.homeDirectory}/.local/share/comfyui";
  venvDir = "${dataDir}/.venv";
  
  startScript = pkgs.writeShellScript "comfyui-start" ''
    set -e
    
    export LD_LIBRARY_PATH="${pkgs.stdenv.cc.cc.lib}/lib:${pkgs.zlib}/lib:${pkgs.libGL}/lib:/run/opengl-driver/lib:$LD_LIBRARY_PATH"
    
    mkdir -p ${dataDir}
    
    # Copy source if not exists or update
    if [ ! -f "${dataDir}/main.py" ]; then
      cp -r ${comfyuiSrc}/* ${dataDir}/
      chmod -R u+w ${dataDir}
    fi
    
    cd ${dataDir}
    
    # Create venv on first run
    if [ ! -d "${venvDir}" ]; then
      ${pkgs.python313}/bin/python -m venv ${venvDir}
      ${venvDir}/bin/pip install --upgrade pip
    fi
    
    # Always ensure deps are installed
    if [ ! -f "${venvDir}/.deps-installed" ]; then
      ${venvDir}/bin/pip install torch torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/cu124
      ${venvDir}/bin/pip install -r requirements.txt
      touch ${venvDir}/.deps-installed
    fi
    
    exec ${venvDir}/bin/python main.py "$@"
  '';
in
{
  home.packages = [ pkgs.python313 pkgs.git ];

  systemd.user.services.comfy-ui = {
    Unit = {
      Description = "ComfyUI";
      After = [ "network.target" ];
    };
    Service = {
      Type = "simple";
      ExecStart = "${startScript}";
      Restart = "on-failure";
      Environment = [ "CUDA_VISIBLE_DEVICES=0" ];
    };
    Install.WantedBy = [ "default.target" ];
  };
}