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

BASE=/home/wildlama/comfy/ComfyUI
MODELS="$BASE/models"
SRC_IMG=/home/wildlama/comfy/zimage-outputs/zimage_latina_00001_.png
INPUT_IMG="$BASE/input/ltx_latina.png"
WF_UI=/home/wildlama/comfy/workflows/ltxv_image_to_video.json
WF_API=/home/wildlama/comfy/workflows/ltx_latina_i2v_api.json
OUT=/home/wildlama/comfy/ltx-outputs
CLEAN_PATH=/home/wildlama/.local/bin:/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin:/snap/bin
export SSL_CERT_FILE=/etc/ssl/certs/ca-certificates.crt
mkdir -p "$MODELS/checkpoints" "$MODELS/text_encoders" "$BASE/input" "$OUT" /home/wildlama/comfy/workflows

log(){ printf '[ltx-robot] %s\n' "$*"; }

download(){
  local url="$1" out="$2" min_bytes="$3"
  if [ -s "$out" ]; then
    local size; size=$(stat -c%s "$out")
    if [ "$size" -ge "$min_bytes" ]; then log "complete: $(basename "$out") ($size bytes)"; return; fi
    log "resume: $(basename "$out") ($size / $min_bytes bytes)"
  else
    log "download: $(basename "$out")"
  fi
  env -u VIRTUAL_ENV PATH="$CLEAN_PATH" curl -L --fail --retry 50 --retry-delay 10 -C - -o "$out" "$url"
  local size; size=$(stat -c%s "$out")
  if [ "$size" -lt "$min_bytes" ]; then log "ERROR too small: $out ($size < $min_bytes)" >&2; exit 1; fi
}

# 1) Models. Use fp8 T5 to keep setup smaller/faster; workflow is patched accordingly.
download 'https://huggingface.co/Lightricks/LTX-Video/resolve/main/ltx-video-2b-v0.9.5.safetensors' "$MODELS/checkpoints/ltx-video-2b-v0.9.5.safetensors" 6300000000
download 'https://huggingface.co/comfyanonymous/flux_text_encoders/resolve/main/t5xxl_fp8_e4m3fn_scaled.safetensors' "$MODELS/text_encoders/t5xxl_fp8_e4m3fn_scaled.safetensors" 4500000000

# 2) Workflow + input image.
if [ ! -s "$WF_UI" ]; then
  cp /home/wildlama/comfy/ComfyUI/.venv/lib/python3.11/site-packages/comfyui_workflow_templates_media_video/templates/ltxv_image_to_video.json "$WF_UI"
fi
cp "$SRC_IMG" "$INPUT_IMG"

# 3) Launch/verify ComfyUI.
log "launch ComfyUI"
env -u VIRTUAL_ENV PATH="$CLEAN_PATH" comfy --workspace "$BASE" launch --background -- --listen 0.0.0.0 --port 8188 || true
for i in $(seq 1 120); do
  if curl -fsS http://127.0.0.1:8188/system_stats >/dev/null 2>&1; then break; fi
  sleep 2
  if [ "$i" = 120 ]; then log "ERROR ComfyUI not ready" >&2; exit 1; fi
done
log "ComfyUI ready"

# 4) Convert official UI template to API and patch prompt/image/model settings.
log "convert workflow"
env -u VIRTUAL_ENV PATH="$CLEAN_PATH" comfy --workspace "$BASE" run --workflow "$WF_UI" --host 127.0.0.1 --port 8188 --print-prompt > /tmp/ltx_prompt_raw.json
python3 - <<'PY'
import json, pathlib, random
raw = pathlib.Path('/tmp/ltx_prompt_raw.json').read_text()
start = raw.find('{')
if start < 0:
    raise SystemExit('No JSON from comfy --print-prompt')
prompt = json.loads(raw[start:])
for node in prompt.values():
    if not isinstance(node, dict):
        continue
    ct = node.get('class_type')
    inp = node.get('inputs', {})
    if ct == 'LoadImage':
        inp['image'] = 'ltx_latina.png'
        if 'upload' in inp: inp['upload'] = 'image'
    elif ct == 'CLIPLoader':
        inp['clip_name'] = 't5xxl_fp8_e4m3fn_scaled.safetensors'
        inp['type'] = 'ltxv'
    elif ct == 'CheckpointLoaderSimple':
        inp['ckpt_name'] = 'ltx-video-2b-v0.9.5.safetensors'
    elif ct == 'CLIPTextEncode':
        old = str(inp.get('text','')).lower()
        if 'low quality' in old or 'worst quality' in old:
            inp['text'] = 'low quality, worst quality, deformed, distorted, disfigured, motion smear, motion artifacts, ugly, bad anatomy'
        else:
            inp['text'] = ('A cinematic portrait video of a beautiful adult Latina woman in a red dress. '
                           'She gently turns her head, blinks naturally, hair moves subtly in a warm breeze, '
                           'soft golden hour light, elegant confident expression, smooth slow camera push-in, realistic, high detail.')
    elif ct == 'LTXVImgToVideo':
        # keep it light and robust locally: 5 seconds at 24 fps-ish template timing
        inp['width'] = 768
        inp['height'] = 512
        inp['length'] = 97
        inp['batch_size'] = 1
        inp['strength'] = 0.12
    elif ct == 'LTXVConditioning':
        inp['frame_rate'] = 24
    elif ct == 'KSamplerSelect':
        inp['sampler_name'] = 'euler'
    elif ct == 'SamplerCustom':
        inp['noise_seed'] = random.randint(1, 2**48-1)
        inp['cfg'] = 3
    elif ct == 'SaveVideo':
        inp['filename_prefix'] = 'ltx_latina'
pathlib.Path('/home/wildlama/comfy/workflows/ltx_latina_i2v_api.json').write_text(json.dumps(prompt, indent=2))
print('wrote api workflow with', len(prompt), 'nodes')
PY

# 5) Run and collect output.
log "run workflow"
RESULT=$(python3 /home/wildlama/.hermes/skills/creative/comfyui/scripts/run_workflow.py --workflow "$WF_API" --output-dir "$OUT" --timeout 1800)
printf '%s\n' "$RESULT" > "$OUT/result.json"
MEDIA=$(python3 - <<'PY'
import json
r=json.load(open('/home/wildlama/comfy/ltx-outputs/result.json'))
outs=r.get('outputs') or []
# Prefer video outputs.
for o in outs:
    if o.get('type') == 'video' or str(o.get('file','')).lower().endswith(('.mp4','.webm','.gif')):
        print(o['file']); break
else:
    print(outs[0]['file'] if outs else '')
PY
)
if [ -z "$MEDIA" ] || [ ! -s "$MEDIA" ]; then log "ERROR no output media" >&2; cat "$OUT/result.json" >&2; exit 1; fi

# 6) Make Telegram/Discord-compatible h264 mp4 when possible.
FINAL="$MEDIA"
if command -v ffmpeg >/dev/null 2>&1 && [[ "$MEDIA" == *.mp4 ]]; then
  FINAL="$OUT/ltx_latina_telegram.mp4"
  ffmpeg -y -i "$MEDIA" -c:v libx264 -profile:v main -preset medium -crf 16 -pix_fmt yuv420p -an "$FINAL" >/dev/null 2>&1 || FINAL="$MEDIA"
fi
log "DONE video: $FINAL"
printf '\nLTX terminé ✅\nMEDIA:%s\n' "$FINAL"