#!/usr/bin/env bash
# gen.sh — 即梦批量出图：提交 → 下载 → 语义化命名 → 记录积分
#
# 用法：
#   bash gen.sh --name <语义名> --prompt-file <提示词文件> [--outdir <目录>]
#                [--ratio 16:9] [--model 5.0] [--res 2k] [--num 1]
#
# 输出：<outdir>/<name>.png（多张为 <name>_1.png …），并打印 submit_id / 扣分。
set -uo pipefail
export PATH="$HOME/.local/bin:$PATH"

NAME=""; PROMPT_FILE=""; OUTDIR="."; RATIO="16:9"; MODEL="5.0"; RES="2k"; NUM=1
while [ $# -gt 0 ]; do
  case "$1" in
    --name)        NAME="$2"; shift 2;;
    --prompt-file) PROMPT_FILE="$2"; shift 2;;
    --outdir)      OUTDIR="$2"; shift 2;;
    --ratio)       RATIO="$2"; shift 2;;
    --model)       MODEL="$2"; shift 2;;
    --res)         RES="$2"; shift 2;;
    --num)         NUM="$2"; shift 2;;
    *) echo "未知参数: $1" >&2; exit 2;;
  esac
done
[ -n "$NAME" ] && [ -n "$PROMPT_FILE" ] || { echo "缺 --name / --prompt-file" >&2; exit 2; }
[ -f "$PROMPT_FILE" ] || { echo "找不到提示词文件: $PROMPT_FILE" >&2; exit 2; }
mkdir -p "$OUTDIR"
PROMPT="$(cat "$PROMPT_FILE")"

CRED_BEFORE=$(dreamina user_credit 2>/dev/null | python3 -c 'import sys,json;print(json.load(sys.stdin)["total_credit"])' 2>/dev/null || echo "?")

OUTLOG=/tmp/gen_$$.out
ERRLOG=/tmp/gen_$$.err
dreamina text2image --prompt="$PROMPT" --ratio="$RATIO" --resolution_type="$RES" \
           --model_version="$MODEL" --generate_num="$NUM" --poll=300 >"$OUTLOG" 2>"$ERRLOG"
SID=$(python3 -c '
import sys,json,re
t=open(sys.argv[1],encoding="utf-8",errors="replace").read()
m=re.search(r"\{.*\}", t, re.S)
try:
    print(json.loads(m.group(0)).get("submit_id","") if m else "")
except Exception:
    print("")
' "$OUTLOG" 2>/dev/null)

if [ -z "$SID" ]; then
  echo "!! 未取到 submit_id，尝试从 list_task 兜底" >&2
  tail -3 "$ERRLOG" >&2
  SID=$(dreamina list_task 2>/dev/null | python3 -c '
import sys,json
try:
    d=json.load(sys.stdin)
    items=d if isinstance(d,list) else (d.get("data") or {}).get("items") or []
    print(items[0].get("submit_id","") if items else "")
except Exception: print("")
' 2>/dev/null)
  if [ -n "$SID" ]; then
    echo "   兜底 submit_id=$SID（注意：可能是上一条任务，请人工核对）" >&2
  else
    echo "!! 放弃，见 $OUTLOG / $ERRLOG" >&2; exit 1
  fi
fi

TMPD=$(mktemp -d)
dreamina query_result --submit_id="$SID" --download_dir="$TMPD" >/dev/null 2>&1
# 轮询直到文件落盘（最多 240s）
for i in $(seq 1 48); do
  n=$(find "$TMPD" -type f \( -name '*.png' -o -name '*.jpg' -o -name '*.jpeg' \) | wc -l)
  [ "$n" -gt 0 ] && break
  sleep 5
  dreamina query_result --submit_id="$SID" --download_dir="$TMPD" >/dev/null 2>&1
done

FILES=$(find "$TMPD" -type f \( -name '*.png' -o -name '*.jpg' -o -name '*.jpeg' \) | sort)
[ -n "$FILES" ] || { echo "!! 未下载到图片 submit_id=$SID" >&2; exit 1; }

idx=0
for f in $FILES; do
  idx=$((idx+1))
  if [ "$idx" -eq 1 ] && [ "$(echo "$FILES" | wc -l)" -eq 1 ]; then
    dest="$OUTDIR/$NAME.png"
  else
    dest="$OUTDIR/${NAME}_${idx}.png"
  fi
  cp "$f" "$dest"
  echo "OK $dest"
done
rm -rf "$TMPD"

CRED_AFTER=$(dreamina user_credit 2>/dev/null | python3 -c 'import sys,json;print(json.load(sys.stdin)["total_credit"])' 2>/dev/null || echo "?")
if [ "$CRED_BEFORE" != "?" ] && [ "$CRED_AFTER" != "?" ]; then
  echo "credit: $CRED_BEFORE -> $CRED_AFTER (扣 $((CRED_BEFORE-CRED_AFTER)))"
else
  echo "credit: $CRED_BEFORE -> $CRED_AFTER"
fi
echo "submit_id: $SID"