#!/usr/bin/env bash
#
# setup.sh — Instala ComfyUI + FLUX.1-dev en tu Mac (Apple Silicon)
# Pensado para macOS con chip M-series. Ejecútalo desde la carpeta Flux2.
#
# Uso:
#   cd ~/Desktop/Flux2
#   chmod +x setup.sh
#   ./setup.sh
#
set -euo pipefail

# ---- Carpeta de instalación = donde está este script ----
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
COMFY_DIR="$SCRIPT_DIR/ComfyUI"

echo "============================================"
echo "  Instalador ComfyUI + FLUX.1-dev"
echo "  Carpeta: $SCRIPT_DIR"
echo "============================================"

# ---- 0. Comprobaciones básicas ----
if [[ "$(uname)" != "Darwin" ]]; then
  echo "AVISO: este script está pensado para macOS. Continuando de todos modos..."
fi
if [[ "$(uname -m)" != "arm64" ]]; then
  echo "AVISO: no se detecta Apple Silicon (arm64). Flux irá MUY lento sin GPU MPS."
fi

# Comprobar git
if ! command -v git >/dev/null 2>&1; then
  echo "ERROR: no tienes 'git'. Instálalo con: xcode-select --install"
  exit 1
fi

# Buscar un Python >= 3.10 (ComfyUI no funciona con 3.9)
PYBIN=""
for cand in python3.12 python3.11 python3.10 python3; do
  if command -v "$cand" >/dev/null 2>&1; then
    ver="$("$cand" -c 'import sys;print("%d.%d"%sys.version_info[:2])' 2>/dev/null || echo 0.0)"
    major="${ver%.*}"; minor="${ver#*.}"
    if [[ "$major" -eq 3 && "$minor" -ge 10 ]]; then
      PYBIN="$cand"
      break
    fi
  fi
done
if [[ -z "$PYBIN" ]]; then
  echo "ERROR: Necesitas Python 3.10 o superior (tienes una versión más antigua)."
  echo "Instálalo con Homebrew:   brew install python@3.12"
  echo "o desde:   https://www.python.org/downloads/"
  exit 1
fi
echo "Python: $($PYBIN --version)  ($PYBIN)"

# ---- 1. Clonar ComfyUI ----
if [[ -d "$COMFY_DIR/.git" ]]; then
  echo "ComfyUI ya existe, actualizando..."
  git -C "$COMFY_DIR" pull --ff-only || true
else
  echo "Clonando ComfyUI..."
  git clone https://github.com/comfyanonymous/ComfyUI.git "$COMFY_DIR"
fi

cd "$COMFY_DIR"

# ---- 2. Entorno virtual ----
# Si ya existe un venv con un Python antiguo (<3.10), lo recreamos.
if [[ -d "venv" ]]; then
  vver="$(venv/bin/python -c 'import sys;print("%d.%d"%sys.version_info[:2])' 2>/dev/null || echo 0.0)"
  vmajor="${vver%.*}"; vminor="${vver#*.}"
  if [[ "$vmajor" -ne 3 || "$vminor" -lt 10 ]]; then
    echo "El venv existente usa Python $vver (demasiado antiguo). Lo recreo con $PYBIN..."
    rm -rf venv
    "$PYBIN" -m venv venv
  fi
else
  echo "Creando entorno virtual (venv) con $PYBIN..."
  "$PYBIN" -m venv venv
fi
# shellcheck disable=SC1091
source venv/bin/activate
python -m pip install --upgrade pip

# ---- 3. PyTorch con soporte MPS (GPU de Apple) ----
echo "Instalando PyTorch (nightly, con MPS)..."
pip install --pre torch torchvision torchaudio \
  --extra-index-url https://download.pytorch.org/whl/nightly/cpu

# ---- 4. Dependencias de ComfyUI ----
echo "Instalando dependencias de ComfyUI..."
pip install -r requirements.txt

# ---- 5. Hugging Face CLI + login ----
pip install -U "huggingface_hub[cli]"

echo ""
echo "--------------------------------------------------------------------"
echo " FLUX.1-dev está protegido (gated). Antes de descargarlo necesitas:"
echo "   1) Tener cuenta en huggingface.co"
echo "   2) Aceptar la licencia en:"
echo "      https://huggingface.co/black-forest-labs/FLUX.1-dev"
echo "   3) Crear un token de acceso (Settings -> Access Tokens)"
echo "--------------------------------------------------------------------"
echo ""

# Login solo si aún no hay token guardado
if ! hf auth whoami >/dev/null 2>&1; then
  echo "Inicia sesión en Hugging Face (pega tu token cuando lo pida):"
  hf auth login
fi

# ---- 6. Descargar los modelos a sus carpetas ----
echo "Descargando modelos (esto puede tardar; son varios GB)..."

# Modelo de difusión
hf download black-forest-labs/FLUX.1-dev flux1-dev.safetensors \
  --local-dir "$COMFY_DIR/models/diffusion_models"

# VAE
hf download black-forest-labs/FLUX.1-dev ae.safetensors \
  --local-dir "$COMFY_DIR/models/vae"

# Text encoders (CLIP-L + T5 XXL fp16, la mejor calidad — perfecto para 128 GB)
hf download comfyanonymous/flux_text_encoders clip_l.safetensors \
  --local-dir "$COMFY_DIR/models/text_encoders"
hf download comfyanonymous/flux_text_encoders t5xxl_fp16.safetensors \
  --local-dir "$COMFY_DIR/models/text_encoders"

# ---- 7. Copiar el workflow de ejemplo ----
if [[ -f "$SCRIPT_DIR/flux_dev_workflow.json" ]]; then
  mkdir -p "$COMFY_DIR/user/default/workflows"
  cp "$SCRIPT_DIR/flux_dev_workflow.json" "$COMFY_DIR/user/default/workflows/" 2>/dev/null || true
fi

echo ""
echo "============================================"
echo "  ¡Listo!"
echo "============================================"
echo "Para arrancar ComfyUI cuando quieras:"
echo ""
echo "   cd \"$COMFY_DIR\""
echo "   source venv/bin/activate"
echo "   python main.py"
echo ""
echo "Luego abre en el navegador:  http://127.0.0.1:8188"
echo "Y carga el workflow:  flux_dev_workflow.json"
echo "============================================"