#!/bin/bash
# Genera tutte le clip di un video sul pod RunPod con ComfyUI, in fila, e le scarica in
# <cartella_video>/clip/. Le clip parlate usano InfiniteTalk (foto + battuta mp3), le scene
# mute Wan 2.2 image-to-video. Non va modificato: legge l'elenco dei lavori da
# <cartella_video>/clip.tsv (vedi clip.tsv.esempio).
#
# Uso: orchestrate.sh <pod_id> <utente_ssh> <cartella_video>
#   pod_id      l'id del pod (es. abc123def456), quello dell'URL https://<pod_id>-8188.proxy.runpod.net
#   utente_ssh  l'utente del proxy ssh (es. abc123def456-64411a2b), lo trovi in "Connect" del pod
#
# Cartella video attesa:
#   clip.tsv            elenco dei lavori
#   immagini/           i fotogrammi di partenza (png/jpg/webp)
#   audio/battute/      le battute mp3, una per clip parlata
#   clip/               dove finiscono le clip (creata dallo script)
#   log/                un log per lavoro (creata dallo script)
#
# Dura da 10 minuti a un'ora: lanciarlo staccato dal terminale, es.
#   nohup ./orchestrate.sh <pod_id> <utente_ssh> <cartella> > <cartella>/orchestrate.log 2>&1 &
# Richiede: bash, curl, ffprobe/ffmpeg, python3, ssh; lib_video.sh, gen_video.py e
# fix_wav2vec_tf5.py nella stessa cartella di questo script.
source "$(dirname "$0")/lib_video.sh" || { echo "manca lib_video.sh accanto a $0: va scaricata insieme agli script"; exit 1; }
D="$(cd "$(dirname "$0")" && pwd)"
POD=$1; SSHUSER=$2; V=$3
G="$D/gen_video.py"; PATCH="$D/fix_wav2vec_tf5.py"
OUT="$V/clip"; LOGS="$V/log"; TSV="$V/clip.tsv"
B=https://$POD-8188.proxy.runpod.net
log(){ _vlog "$@"; }
fail(){ log "ORCH_FAILED: $*"; exit 1; }
[ -n "$POD" ] && [ -n "$SSHUSER" ] && [ -d "$V" ] || fail "uso: orchestrate.sh <pod_id> <utente_ssh> <cartella_video>"
[ -f "$TSV" ] || fail "manca $TSV (vedi clip.tsv.esempio)"
[ -f "$G" ] && [ -f "$PATCH" ] || fail "mancano gen_video.py o fix_wav2vec_tf5.py accanto allo script"

# ---- 1. Legge clip.tsv: nome, tipo (it|wan), immagine, audio (it) o fotogrammi (wan), prompt
declare -a NOMI TIPI IMGS AUD PROMPTS
n_it=0; n_wan=0
while IFS=$'\t' read -r nome tipo img aud prompt; do
  [ -z "$nome" ] || [ "${nome:0:1}" = "#" ] && continue
  case "$tipo" in
    it)  [ -f "$V/audio/battute/$aud" ] || fail "riga '$nome': battuta $V/audio/battute/$aud non trovata"; n_it=$((n_it+1)) ;;
    wan) [[ "$aud" =~ ^[0-9]+$ ]] || fail "riga '$nome': per wan la quarta colonna è il numero di fotogrammi (es. 81 = 5 s a 16 fps)"; n_wan=$((n_wan+1)) ;;
    *)   fail "riga '$nome': tipo '$tipo' sconosciuto (it oppure wan)" ;;
  esac
  [ -f "$V/immagini/$img" ] || fail "riga '$nome': immagine $V/immagini/$img non trovata"
  [ -n "$prompt" ] || fail "riga '$nome': prompt vuoto"
  NOMI+=("$nome"); TIPI+=("$tipo"); IMGS+=("$img"); AUD+=("$aud"); PROMPTS+=("$prompt")
done < "$TSV"
[ ${#NOMI[@]} -gt 0 ] || fail "clip.tsv non ha righe valide"
log "LAVORI ${#NOMI[@]} (parlate: $n_it, mute: $n_wan)"

# ---- 2. Prima di accendere la GPU: una battuta muta genera una clip col labiale sbagliato e va rifatta
for i in "${!NOMI[@]}"; do
  [ "${TIPI[$i]}" = it ] && { controlla_audio "$V/audio/battute/${AUD[$i]}" 1.5 || fail "battuta muta: ${AUD[$i]}"; }
done

comfy_attendi "$POD" "$SSHUSER" 1200 || fail "ComfyUI non parte"

# ---- 3. Aspetta che i modelli siano sul pod: ComfyUI risponde prima che il template abbia finito di scaricarli
need_wan="wan2.2_i2v_high_noise_14B_fp16.safetensors wan2.2_i2v_low_noise_14B_fp16.safetensors i2v_lightx2v_high_noise_model.safetensors i2v_lightx2v_low_noise_model.safetensors umt5_xxl_fp8_e4m3fn_scaled.safetensors wan_2.1_vae.safetensors"
# wav2vec2 non si aspetta qui: ComfyUI lo vede solo dopo il collegamento del passo 4, che poi
# ne controlla la presenza (w2v=1). Aspettarlo qui bloccava lo script 30 minuti e lo faceva fallire.
need_it="Wan2_1-InfiniTetalk-Single_fp16.safetensors wan2.1_i2v_480p_14B_bf16.safetensors lightx2v_I2V_14B_480p_cfg_step_distill_rank64_bf16.safetensors umt5-xxl-enc-bf16.safetensors Wan2_1_VAE_bf16.safetensors"
need="$([ $n_wan -gt 0 ] && echo "$need_wan") $([ $n_it -gt 0 ] && echo "$need_it")"
t0=$(date +%s)
while :; do
  have=$(for d in diffusion_models loras text_encoders vae; do curl -s -A "$UA" --max-time 20 "$B/models/$d"; echo; done)
  miss=""; for m in $need; do echo "$have" | grep -q "$m" || miss="$miss $m"; done
  [ -z "$miss" ] && break
  if [ $(( $(date +%s)-t0 )) -gt 1800 ]; then
    fail "dopo 30 minuti mancano ancora questi modelli:$miss (i file di InfiniteTalk vanno aggiunti a mano sul disco: vedi la guida)"
  fi
  sleep 30
done
log "MODELS tutti presenti dopo $(( $(date +%s)-t0 ))s"
[ $(( $(date +%s)-t0 )) -gt 5 ] && { log "aspetto 90s che i download si assestino"; sleep 90; }

# ---- 4. Solo con clip parlate: collegamento wav2vec2, patch per transformers 5, riavvio di ComfyUI.
#         Sono sul disco del container, si perdono a ogni pod nuovo: si rifanno ogni volta.
if [ $n_it -gt 0 ]; then
  # La patch viaggia a righe da 400 caratteri e si scrive nel /tmp del container: una riga
  # sola lunghissima dentro la sessione ssh è fragile, e su un disco di rete pieno la
  # copia in /workspace arrivava vuota (Disk quota exceeded) senza errori visibili.
  patch_lines=$(base64 -w 400 "$PATCH" | sed "s|.*|echo '&' >> /tmp/fix_wav2vec_tf5.b64|")
  out=$(ssh_pod "$SSHUSER" "$(printf '%s\n' \
   'rm -f /tmp/fix_wav2vec_tf5.b64' \
   "$patch_lines" \
   'base64 -d /tmp/fix_wav2vec_tf5.b64 > /tmp/fix_wav2vec_tf5.py' \
   '[ -e /ComfyUI/models/wav2vec2 ] || ln -s /workspace/ComfyUI/models/wav2vec2 /ComfyUI/models/wav2vec2' \
   'python3 /tmp/fix_wav2vec_tf5.py 2>&1 | sed "s/^/Z$((1+1))P /"' \
   'echo "Z$((1+1))M markers=$(grep -c "PATCH transformers>=5 hidden_states" /ComfyUI/custom_nodes/ComfyUI-WanVideoWrapper/multitalk/wav2vec2.py) w2v=$(ls /ComfyUI/models/wav2vec2/ | wc -l)"' \
   "pkill -f 'ComfyUI/main.py'; sleep 5; cd /ComfyUI && setsid nohup python3 /ComfyUI/main.py --listen 0.0.0.0 --enable-cors-header '*' --use-sage-attention --extra-model-paths-config /ComfyUI/extra_model_paths.yaml >> /workspace/comfyui_manual.log 2>&1 < /dev/null & disown; sleep 3; echo \"Z\$((1+1))R procs=\$(pgrep -fc 'ComfyUI/main.py')\"")" \
   | grep -o -E 'Z2[PMR] .*' | grep -v -E '\$\(\(|\\\$')
  log "PATCH output: $(echo "$out" | tr '\n' ' ')"
  echo "$out" | grep -q 'Z2M markers=2 w2v=1' || fail "patch o collegamento wav2vec2 mancanti (vedi la guida, sezione InfiniteTalk)"
  echo "$out" | grep -q 'Z2R procs=[1-9]' || fail "ComfyUI non è ripartito dopo la patch"
  sleep 15
  comfy_attendi "$POD" "$SSHUSER" 300 || fail "ComfyUI non risponde dopo il riavvio"
fi

# ---- 5. Carica immagini e battute (una volta sola ciascuna)
declare -A caricati
carica(){
  local f=$1 r
  [ -n "${caricati[$f]}" ] && return 0
  r=$(curl -s -A "$UA" --max-time 120 -F "image=@$f" -F "type=input" -F "overwrite=true" "$B/upload/image")
  echo "$r" | grep -q '"name"' || fail "upload fallito per $f: $r"
  caricati[$f]=1
}
for i in "${!NOMI[@]}"; do
  carica "$V/immagini/${IMGS[$i]}"
  [ "${TIPI[$i]}" = it ] && carica "$V/audio/battute/${AUD[$i]}"
done
log "UPLOADED ${#caricati[@]} file"

# ---- 6. Un lavoro alla volta, con log separato e conteggio finale
dur(){ ffprobe -v error -show_entries format=duration -of csv=p=0 "$1"; }
ok=0; ko=0
mkdir -p "$OUT" "$LOGS"
for i in "${!NOMI[@]}"; do
  nome=${NOMI[$i]}
  if [ "${TIPI[$i]}" = it ]; then
    python3 "$G" it "$B" "${IMGS[$i]}" "${AUD[$i]}" "${PROMPTS[$i]}" "$nome" "$OUT" "$(dur "$V/audio/battute/${AUD[$i]}")" > "$LOGS/job_$nome.log" 2>&1; rc=$?
  else
    python3 "$G" wan "$B" "${IMGS[$i]}" "${AUD[$i]}" "${PROMPTS[$i]}" "$nome" "$OUT" > "$LOGS/job_$nome.log" 2>&1; rc=$?
  fi
  while IFS= read -r l; do log "$l"; done < "$LOGS/job_$nome.log"
  if [ $rc = 0 ] && [ -s "$OUT/$nome.mp4" ]; then ok=$((ok+1)); else ko=$((ko+1)); log "JOB_FAILED $nome rc=$rc"; fi
done
log "ALL_DONE ok=$ok failed=$ko"
[ $ko = 0 ]