#!/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 / 扣分。
#
# 防重护栏（重要）：
#   提交前记录当前最新 submit_id（PREV_SID）。若提交阶段网络超时导致拿不到
#   submit_id，先重试提交；仍失败才用 list_task 兜底，且兜底值必须 != PREV_SID，
#   否则判定为「未真正提交」，直接失败退出——避免把上一条任务的图当成本次的产物。
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")"

extract_sid() { # extract_sid <文件>
  # ⚠️ 不能用贪婪正则 \{.*\}：即梦的 SLOW SQL 日志行里含有
  #    commerce_info="{...}"，从那里的 { 开始匹配会解析失败。
  #    正确做法是从「行首单独一个 {」开始的最后一段 JSON 解析。
  python3 -c '
import sys,json
t=open(sys.argv[1],encoding="utf-8",errors="replace").read()
i=t.rfind("\n{\n")
seg=t[i+1:] if i>=0 else t
try:
    print(json.loads(seg).get("submit_id",""))
except Exception:
    print("")
' "$1" 2>/dev/null
}

newest_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
}

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

CRED_BEFORE=$(cred)
PREV_SID=$(newest_sid)          # 提交前的最新任务，用于防重
echo "    [护栏] 提交前最新 submit_id = ${PREV_SID:-<空>}"

TAG=/tmp/gen_$$
SID=""
for attempt in 1 2 3; do
  dreamina text2image --prompt="$PROMPT" --ratio="$RATIO" --resolution_type="$RES" \
           --model_version="$MODEL" --generate_num="$NUM" --poll=300 \
           >"${TAG}.out" 2>"${TAG}.err"
  SID=$(extract_sid "${TAG}.out")
  [ -n "$SID" ] && break
  echo "    [重试 $attempt/3] 本次未取到 submit_id：$(tail -1 "${TAG}.err" 2>/dev/null | cut -c1-120)" >&2
  sleep 5
done

if [ -z "$SID" ]; then
  NEW_SID=$(newest_sid)
  if [ -n "$NEW_SID" ] && [ "$NEW_SID" != "$PREV_SID" ]; then
    SID="$NEW_SID"
    echo "    [兜底] 采用新出现的 submit_id=$SID（已确认不等于提交前的最新任务）" >&2
  else
    echo "!! 提交失败：list_task 最新任务仍是提交前的 ${PREV_SID:-<空>}，说明本次并未真正提交。" >&2
    echo "   拒绝使用该任务，以免把上一条的图当成本次产物。见 ${TAG}.err" >&2
    exit 1
  fi
fi

TMPD=$(mktemp -d)
for i in $(seq 1 48); do
  dreamina query_result --submit_id="$SID" --download_dir="$TMPD" >/dev/null 2>&1
  n=$(find "$TMPD" -type f \( -name '*.png' -o -name '*.jpg' -o -name '*.jpeg' \) | wc -l)
  [ "$n" -gt 0 ] && break
  sleep 5
done

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

TOTAL=$(echo "$FILES" | wc -l)
idx=0
for f in $FILES; do
  idx=$((idx+1))
  if [ "$TOTAL" -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=$(cred)
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"