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

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

  controlnetAux = pkgs.fetchFromGitHub {
    owner = "Fannovel16";
    repo = "comfyui_controlnet_aux";
    rev = "136f125c89aed92ced1b6fbb491e13719b72fcc0";
    sha256 = "DlspkqzN7Ls8kXWQMtVQygzsgu/z6FtjMqDthuza/Kc=";
  };

  ipAdapter = pkgs.fetchFromGitHub {
    owner = "cubiq";
    repo = "ComfyUI_IPAdapter_plus";
    rev = "main";
    sha256 = "Ft9WJcmjzon2tAMJq5na24iqYTnQWEQFSKUElSVwYgw=";
  };

  toolingNodes = pkgs.fetchFromGitHub {
    owner = "Acly";
    repo = "comfyui-tooling-nodes";
    rev = "main";
    sha256 = "tVvpVWDpihy7zdV/L7cOpsWE68l15xKIwuM3EriUM+Y=";
  };

  inpaintNodes = pkgs.fetchFromGitHub {
    owner = "Acly";
    repo = "comfyui-inpaint-nodes";
    rev = "main";
    sha256 = "Uy6ppXNAQAOIkmoJB8miAzVUXZ0Elyp+w+kwNxWZjvo=";
  };

  dataDir = "${config.home.homeDirectory}/.local/share/comfyui";
  venvDir = "${dataDir}/.venv";
  customNodesDir = "${dataDir}/custom_nodes";
  kritaModelsDir = "${config.home.homeDirectory}/.var/app/org.kde.krita/data/krita/ai_diffusion/server/models";
  
  startScript = pkgs.writeShellScript "comfyui-start" ''
    set -e
    
    export LD_LIBRARY_PATH="${pkgs.stdenv.cc.cc.lib}/lib:${pkgs.zlib}/lib:${pkgs.libGL}/lib:${pkgs.glib}/lib:${pkgs.xorg.libxcb}/lib:${pkgs.xorg.libX11}/lib:${pkgs.xorg.libXext}/lib:/run/opengl-driver/lib:$LD_LIBRARY_PATH"
    
    mkdir -p ${dataDir}
    mkdir -p ${customNodesDir}
    
    # Copy source if not exists
    if [ ! -f "${dataDir}/main.py" ]; then
      cp -r ${comfyuiSrc}/* ${dataDir}/
      chmod -R u+w ${dataDir}
    fi
    
    # Create extra_model_paths.yaml to use Krita AI Diffusion models
    cat > ${dataDir}/extra_model_paths.yaml << EOF
krita_ai:
    base_path: ${kritaModelsDir}
    checkpoints: checkpoints/
    clip_vision: clip_vision/
    controlnet: controlnet/
    diffusion_models: diffusion_models/
    embeddings: embeddings/
    inpaint: inpaint/
    ipadapter: ipadapter/
    loras: loras/
    style_models: style_models/
    text_encoders: text_encoders/
    upscale_models: upscale_models/
    vae: vae/
EOF
    
    # Install custom nodes
    if [ ! -d "${customNodesDir}/comfyui_controlnet_aux" ]; then
      cp -r ${controlnetAux} ${customNodesDir}/comfyui_controlnet_aux
      chmod -R u+w ${customNodesDir}/comfyui_controlnet_aux
    fi
    
    if [ ! -d "${customNodesDir}/ComfyUI_IPAdapter_plus" ]; then
      cp -r ${ipAdapter} ${customNodesDir}/ComfyUI_IPAdapter_plus
      chmod -R u+w ${customNodesDir}/ComfyUI_IPAdapter_plus
    fi
    
    if [ ! -d "${customNodesDir}/comfyui-tooling-nodes" ]; then
      cp -r ${toolingNodes} ${customNodesDir}/comfyui-tooling-nodes
      chmod -R u+w ${customNodesDir}/comfyui-tooling-nodes
    fi
    
    if [ ! -d "${customNodesDir}/comfyui-inpaint-nodes" ]; then
      cp -r ${inpaintNodes} ${customNodesDir}/comfyui-inpaint-nodes
      chmod -R u+w ${customNodesDir}/comfyui-inpaint-nodes
    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
    
    # Install deps
    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
    
    # Install custom nodes deps
    if [ ! -f "${venvDir}/.custom-nodes-installed" ]; then
      for node in ${customNodesDir}/*/; do
        if [ -f "$node/requirements.txt" ]; then
          # Replace opencv-python with headless version (no GUI deps)
          sed 's/opencv-python>=/opencv-python-headless>=/g' "$node/requirements.txt" > /tmp/requirements_patched.txt
          ${venvDir}/bin/pip install -r /tmp/requirements_patched.txt || true
        fi
      done
      touch ${venvDir}/.custom-nodes-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" ];
  };
}