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

BASE=/home/wildlama/comfy/ComfyUI
MODELS="$BASE/models"
WF_UI=/home/wildlama/comfy/workflows/z_image_turbo_bf16-DF11-workflow.json
WF_API=/home/wildlama/comfy/workflows/z_image_turbo_latina_api.json
OUT=/home/wildlama/comfy/zimage-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/diffusion_models" "$MODELS/text_encoders" "$MODELS/vae" "$OUT" /home/wildlama/comfy/workflows

log(){ printf '[zimage-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) Dependencies/models.
if [ ! -d "$BASE/custom_nodes/ComfyUI-DFloat11-Extended" ]; then
  log "install DF11 custom node"
  env -u VIRTUAL_ENV PATH="$CLEAN_PATH" comfy --workspace "$BASE" stop || true
  env -u VIRTUAL_ENV PATH="$CLEAN_PATH" comfy --workspace "$BASE" node install https://github.com/mingyi456/ComfyUI-DFloat11-Extended
fi

download 'https://huggingface.co/mingyi456/Z-Image-Turbo-DF11-ComfyUI/resolve/main/z_image_turbo_bf16-DF11.safetensors' "$MODELS/diffusion_models/z_image_turbo_bf16-DF11.safetensors" 8300000000
download 'https://huggingface.co/tsqn/Z-Image-Turbo_fp8_comfyui/resolve/main/qwen_3_4b_bf16_fp8_scaled.safetensors' "$MODELS/text_encoders/qwen_3_4b_bf16_fp8_scaled.safetensors" 4400000000
download 'https://huggingface.co/tsqn/Z-Image-Turbo_fp32-fp16-bf16_comfyui/resolve/main/vae/ae.safetensors' "$MODELS/vae/ae.safetensors" 330000000

# Workflow expects qwen_3_4b.safetensors; point that name at the downloaded quantized encoder if needed.
if [ ! -e "$MODELS/text_encoders/qwen_3_4b.safetensors" ]; then
  ln -s qwen_3_4b_bf16_fp8_scaled.safetensors "$MODELS/text_encoders/qwen_3_4b.safetensors"
fi

if [ ! -s "$WF_UI" ]; then
  log "download workflow"
  env -u VIRTUAL_ENV PATH="$CLEAN_PATH" curl -L --fail -o "$WF_UI" https://huggingface.co/mingyi456/Z-Image-Turbo-DF11-ComfyUI/resolve/main/z_image_turbo_bf16-DF11-workflow.json
fi

# 2) Launch ComfyUI.
log "launch ComfyUI"
env -u VIRTUAL_ENV PATH="$CLEAN_PATH" comfy --workspace "$BASE" stop || true
env -u VIRTUAL_ENV PATH="$CLEAN_PATH" comfy --workspace "$BASE" launch --background -- --listen 0.0.0.0 --port 8188
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"

# 3) Convert official UI workflow to API format using comfy-cli, then customize prompt.
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/zimage_prompt_raw.json
python3 - <<'PY'
import json, pathlib, random
raw = pathlib.Path('/tmp/zimage_prompt_raw.json').read_text()
start = raw.find('{')
if start < 0:
    raise SystemExit('No JSON object in comfy --print-prompt output')
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 == 'CLIPTextEncode':
        old = str(inp.get('text','')).lower()
        if 'blurry' in old or 'bad' in old or 'ugly' in old:
            inp['text'] = 'blurry, ugly, bad anatomy, low quality, deformed, extra fingers'
        else:
            inp['text'] = ('beautiful adult Latina woman, elegant portrait photo, warm brown skin, '
                           'dark wavy hair, confident expression, cinematic golden hour lighting, '
                           'stylish red dress, realistic, high detail, 85mm lens, shallow depth of field')
    elif ct == 'KSampler':
        inp['seed'] = random.randint(1, 2**48-1)
        inp['steps'] = 9
        inp['cfg'] = 1
        inp['sampler_name'] = 'euler'
        inp['scheduler'] = 'simple'
        inp['denoise'] = 1
    elif ct == 'EmptySD3LatentImage':
        inp['width'] = 1024
        inp['height'] = 1024
        inp['batch_size'] = 1
    elif ct == 'SaveImage':
        inp['filename_prefix'] = 'zimage_latina'
pathlib.Path('/home/wildlama/comfy/workflows/z_image_turbo_latina_api.json').write_text(json.dumps(prompt, indent=2))
print('wrote api workflow with', len(prompt), 'nodes')
PY

# 4) Execute via skill runner so outputs are copied to OUT.
log "run workflow"
RESULT=$(python3 /home/wildlama/.hermes/skills/creative/comfyui/scripts/run_workflow.py --workflow "$WF_API" --output-dir "$OUT" --timeout 1200)
printf '%s\n' "$RESULT" > "$OUT/result.json"
IMG=$(python3 - <<'PY'
import json
r=json.load(open('/home/wildlama/comfy/zimage-outputs/result.json'))
outs=r.get('outputs') or []
print(outs[0]['file'] if outs else '')
PY
)
if [ -z "$IMG" ] || [ ! -s "$IMG" ]; then log "ERROR no output image" >&2; cat "$OUT/result.json" >&2; exit 1; fi
log "DONE image: $IMG"
printf '\nZ-Image Turbo terminé ✅\nMEDIA:%s\n' "$IMG"