#!/usr/bin/env bash
set -euo pipefail
cd "$(dirname "$0")"

export UV_CACHE_DIR="$(pwd)/.uv-cache"
export UV_SKIP_WHEEL_FILENAME_CHECK=1
mkdir -p "$UV_CACHE_DIR"

# --- 1) ensure uv ---------------------------------------------------------
if ! command -v uv >/dev/null 2>&1; then
    echo "[QwenImageEditLite] 'uv' not found on PATH. Installing uv ..."
    if command -v curl >/dev/null 2>&1; then
        curl -LsSf https://astral.sh/uv/install.sh | sh
    elif command -v wget >/dev/null 2>&1; then
        wget -qO- https://astral.sh/uv/install.sh | sh
    else
        echo "[QwenImageEditLite] need curl or wget to install uv. Install one and retry,"
        echo "                     or install uv manually: https://docs.astral.sh/uv/"
        exit 1
    fi
    export PATH="$HOME/.local/bin:$HOME/.cargo/bin:$PATH"
fi

# --- 2) platform hint ----------------------------------------------------
case "$(uname -s)" in
    Linux)
        case "$(uname -m)" in
            aarch64)
                PLATFORM_KIND="dgx-spark"
                echo "[QwenImageEditLite] Linux aarch64 detected (DGX Spark path)."
                echo "[QwenImageEditLite] Creating Python 3.13 virtual environment."
                echo "[QwenImageEditLite] Syncing dependencies with CUDA 13.0 PyTorch."
                echo "[QwenImageEditLite] Allowing known wheel filename/version mismatch for bundled custom wheels."
                ;;
            x86_64)
                PLATFORM_KIND="linux-x64"
                echo "[QwenImageEditLite] Linux x86_64 detected."
                echo "[QwenImageEditLite] Creating Python 3.13 virtual environment."
                echo "[QwenImageEditLite] Syncing dependencies with CUDA 13.0 PyTorch."
                echo "[QwenImageEditLite] Allowing known wheel filename/version mismatch for bundled custom wheels."
                ;;
            *)
                echo "[QwenImageEditLite] Unsupported Linux architecture: $(uname -m)"
                exit 1
                ;;
        esac
        ;;
    Darwin)
        echo "[QwenImageEditLite] macOS is not a target for this project."
        echo "[QwenImageEditLite] Supported targets are Windows x64, Linux x64, and DGX Spark."
        exit 1
        ;;
    *)
        echo "[QwenImageEditLite] Unsupported OS: $(uname -s)"
        exit 1
        ;;
esac

# --- 3) create venv ------------------------------------------------------
if [ -d ".venv" ]; then
    echo "[QwenImageEditLite] Existing .venv detected."
    if [ ! -t 0 ]; then
        echo "[QwenImageEditLite] Install cancelled because confirmation is required to delete .venv."
        exit 1
    fi
    read -r -p "[QwenImageEditLite] Delete the existing .venv and continue? [y/N] " reply
    case "$reply" in
        [yY]|[yY][eE][sS])
            echo "[QwenImageEditLite] Removing existing .venv ..."
            rm -rf .venv
            ;;
        *)
            echo "[QwenImageEditLite] Install cancelled. Existing .venv was kept."
            exit 1
            ;;
    esac
fi

uv venv .venv --python 3.13

# --- 4) install deps -----------------------------------------------------
uv sync --extra cu130

# --- 5) verify bundled runtime -------------------------------------------
if [ ! -f "vendor/comfyui/nodes.py" ]; then
    echo "[QwenImageEditLite] bundled vendor/comfyui runtime not found."
    exit 1
fi

if [ ! -f "config.ini" ]; then
    echo "[QwenImageEditLite] config.ini not found."
    exit 1
fi

cat <<'EOF'

[QwenImageEditLite] Install complete.
[QwenImageEditLite] Virtual environment:
  .venv (Python 3.13)
[QwenImageEditLite] Bundled runtime:
  vendor/comfyui
[QwenImageEditLite] Attention backend:
  optional CUDA attention packages synced from pyproject.toml
[QwenImageEditLite] Put model files under:
  models/diffusion_models
  models/vae
  models/text_encoders
[QwenImageEditLite] Put LoRA files under:
  loras
[QwenImageEditLite] Input images can be placed under:
  input
EOF