#!/usr/bin/env bash
set -euo pipefail

echo "======================================================================"
echo " ADUC-SDR Video Suite — Inicialização (SeedVR2-3B + FastAPI/Gradio UI)"
echo "======================================================================"

# --- 1) Ambiente e GPUs -------------------------------------------------------
export PYTHONFAULTHANDLER=1
export TORCH_NCCL_ASYNC_ERROR_HANDLING=1
export NCCL_DEBUG="${NCCL_DEBUG:-INFO}"
export TORCH_DISTRIBUTED_DEBUG="${TORCH_DISTRIBUTED_DEBUG:-DETAIL}"

# Autodetecta GPUs somente se a variável não estiver definida
if [ -z "${CUDA_VISIBLE_DEVICES:-}" ]; then
  if command -v nvidia-smi >/dev/null 2>&1; then
    export CUDA_VISIBLE_DEVICES="$(nvidia-smi --query-gpu=index --format=csv,noheader | tr '\n' ',' | sed 's/,$//')"
  fi
fi
NUM_GPUS=$(python3 - <<'PY'
import os, torch
print(torch.cuda.device_count() if torch.cuda.is_available() else 0)
PY
)
echo "[INFO] GPUs visíveis: ${NUM_GPUS} (${CUDA_VISIBLE_DEVICES:-unset})"

# --- 2) Builder CUDA (opcional) -----------------------------------------------
echo "🛠️ [ETAPA 1/6] Executando builder.sh (se existir)..."
if [ -f "/app/builder.sh" ]; then
  bash /app/builder.sh
  echo "✅ Builder finalizado."
else
  echo "⚠️ Aviso: builder.sh não encontrado. Pulando etapa."
fi

# --- 3) Caches persistentes (HF/Torch) ----------------------------------------
if [ -d /data ]; then
  echo "[INFO] Usando /data como raiz de cache persistente."
  export HF_HOME="${HF_HOME:-/data/.cache/huggingface}"
  export TORCH_HOME="${TORCH_HOME:-/data/.cache/torch}"
else
  echo "[INFO] Usando /app/.cache como fallback local."
  export HF_HOME="${HF_HOME:-/app/.cache/huggingface}"
  export TORCH_HOME="${TORCH_HOME:-/app/.cache/torch}"
fi
export HF_HUB_CACHE="${HF_HUB_CACHE:-$HF_HOME/hub}"
export HF_HUB_ENABLE_HF_TRANSFER="${HF_HUB_ENABLE_HF_TRANSFER:-1}"
mkdir -p "$HF_HUB_CACHE" "$TORCH_HOME"
echo "[INFO] HF_HOME=$HF_HOME"
echo "[INFO] TORCH_HOME=$TORCH_HOME"

# --- 4) SeedVR: código-fonte e checkpoints ------------------------------------
export MODEL_REPO="${MODEL_REPO:-ByteDance-Seed/SeedVR2-3B}"
export CKPT_DIR="${CKPT_DIR:-/app/ckpts}"
export SRC_REPO_DIR="${SRC_REPO_DIR:-/app/SeedVR_source}"

echo "📦 [ETAPA 2/6] Código-fonte SeedVR..."
if [ ! -d "$SRC_REPO_DIR/.git" ] && [ ! -d "$SRC_REPO_DIR/projects" ]; then
  git clone https://github.com/ByteDance-Seed/SeedVR.git "$SRC_REPO_DIR" || true
else
  echo "[INFO] SeedVR já presente em $SRC_REPO_DIR."
fi

#echo "📁 [ETAPA 3/6] Sincronizando módulos em /app..."
#mkdir -p /app/common /app/projects /app/data /app/models /app/configs_3b "$CKPT_DIR"
# Copia se existirem; cria symlink se faltarem
#for d in common projects data models configs_3b; do
#  if [ -d "$SRC_REPO_DIR/$d" ]; then
#    cp -rvu "$SRC_REPO_DIR/$d/." "/app/$d/." || true
#  fi
#  if [ ! -d "/app/$d" ] && [ -d "$SRC_REPO_DIR/$d" ]; then
#    #ln -s "$SRC_REPO_DIR/$d" "/app/$d"
#    echo "[LINK] /app/$d -> $SRC_REPO_DIR/$d"
#  fi
#done

# Pré-download dos pesos e embeddings
echo "📥 [ETAPA 4/6] Verificando/baixando pesos do modelo..."
python3 - <<'PY'
import os, sys, traceback
from huggingface_hub import snapshot_download
try:
    repo_id   = os.environ.get("MODEL_REPO", "ByteDance-Seed/SeedVR2-3B")
    cache_dir = os.environ.get("HF_HUB_CACHE")
    ckpt_dir  = os.environ.get("CKPT_DIR", "/app/ckpts")
    os.makedirs(ckpt_dir, exist_ok=True)
    print(f"[HF] snapshot {repo_id} -> {ckpt_dir}")
    snapshot_download(
        repo_id=repo_id,
        cache_dir=cache_dir,
        local_dir=ckpt_dir,
        #local_dir_use_symlinks=False,
        #force_download=True,
    )
    from torch.hub import download_url_to_file
    for name in ("pos_emb.pt","neg_emb.pt"):
        dst = f"/app/{name}"
        if not os.path.exists(dst):
            url = f"https://huggingface.co/{repo_id}/resolve/main/{name}"
            print(f"[DL] {url} -> {dst}")
            download_url_to_file(url, dst)
    print("✅ Pesos e embeddings prontos.")
except Exception as e:
    print(f"🛑 Falha no prefetch: {e}")
    traceback.print_exc()
PY

# Garante que os imports absolutos funcionem
export PYTHONPATH="/app:/app/projects:${SRC_REPO_DIR}:${PYTHONPATH:-}"

# --- 5) VINCIE persistente em /data/VINCIE (ou /app/VINCIE) -------------------
if [ -d /data ]; then PERSIST_ROOT="/data"; else PERSIST_ROOT="/app"; fi
export VINCIE_DIR="${VINCIE_DIR:-${PERSIST_ROOT}/VINCIE}"
mkdir -p "$VINCIE_DIR"
echo "[VINCIE] destino: $VINCIE_DIR"

# Defina um dos modos abaixo via env:
# export VINCIE_REMOTE="https://github.com/<ORG>/<REPO>.git"  # Git
# export VINCIE_REF="main"
# export VINCIE_HF_REPO="org/VINCIE"                          # Hugging Face

# Git mode
if [ -n "${VINCIE_REMOTE:-}" ]; then
  if [ -d "$VINCIE_DIR/.git" ]; then
    echo "[VINCIE] Atualizando repo Git..."
    git -C "$VINCIE_DIR" fetch --all -p --tags
    git -C "$VINCIE_DIR" checkout "${VINCIE_REF:-main}"
    git -C "$VINCIE_DIR" pull --ff-only
  else
    echo "[VINCIE] Clonando ${VINCIE_REMOTE} -> ${VINCIE_DIR}"
    git clone --depth 1 --branch "${VINCIE_REF:-main}" "$VINCIE_REMOTE" "$VINCIE_DIR"
  fi
fi

# HF mode
if [ -n "${VINCIE_HF_REPO:-}" ]; then
  python3 - <<'PY'
import os
from huggingface_hub import snapshot_download
repo_id   = os.environ["VINCIE_HF_REPO"]
local_dir = os.environ["VINCIE_DIR"]
cache_dir = os.environ.get("HF_HUB_CACHE")
os.makedirs(local_dir, exist_ok=True)
snapshot_download(
    repo_id=repo_id,
    cache_dir=cache_dir,
    local_dir=local_dir,
)
print(f"[VINCIE] snapshot pronto em {local_dir}")
PY
fi

# Torna VINCIE importável, se contiver pacotes Python
export PYTHONPATH="${VINCIE_DIR}:${PYTHONPATH:-}"

# --- Pretty tree (somente dirs, limpo) ---
echo "🌳 Estrutura"




tree -L 7 /app
echo "  ############## 
"

tree -L 6 /data
echo "  ############## 
"




# --- 7) Lançamento da API (Uvicorn) -------------------------------------------
echo "🚀 Iniciando app.py via Uvicorn..."
APP_HOST="${APP_HOST:-0.0.0.0}"
APP_PORT="${APP_PORT:-8000}"
APP_WORKERS="${APP_WORKERS:-1}"
# app.py deve exportar 'app = FastAPI(...)'
exec python -m uvicorn app:app --host "${APP_HOST}" --port "${APP_PORT}" --workers "${APP_WORKERS}" --proxy-headers