#!/usr/bin/env bash
# img2img.sh — 即梦图生图（可靠版）：提交 → 取 submit_id（stdout 优先，tasks.db 按精确 Prompt 反查兜底）→ 下载
#
# 用法：
#   bash img2img.sh --name <语义名> --prompt-file <提示词> --outdir <目录> \
#        --image <输入图1> [--image <输入图2> ...] [--ratio 16:9|3:4] [--model 5.0] [--res 2k] [--num 1]
set -uo pipefail
export PATH="$HOME/.local/bin:$PATH"

NAME=""; PROMPT_FILE=""; OUTDIR="."; RATIO=""; MODEL="5.0"; RES="2k"; NUM=1; IMAGES=()
while [ $# -gt 0 ]; do
  case "$1" in
    --name)        NAME="$2"; shift 2;;
    --prompt-file) PROMPT_FILE="$2"; shift 2;;
    --outdir)      OUTDIR="$2"; shift 2;;
    --image)       IMAGES+=("$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; }
[ ${#IMAGES[@]} -gt 0 ] || { echo "缺 --image（至少 1 张输入图）" >&2; exit 2; }
for f in "${IMAGES[@]}"; do [ -f "$f" ] || { echo "找不到输入图: $f" >&2; exit 2; }; done
mkdir -p "$OUTDIR"
PROMPT="$(cat "$PROMPT_FILE")"

credit() { dreamina user_credit 2>/dev/null | python3 -c 'import sys,json,re
t=sys.stdin.read(); m=re.search(r"\{.*\}",t,re.S)
print(json.loads(m.group(0))["total_credit"] if m else "?")' 2>/dev/null || echo "?"; }

ARGS=(--prompt="$PROMPT" --resolution_type="$RES" --model_version="$MODEL" --generate_num="$NUM" --poll=1)
[ -n "$RATIO" ] && ARGS+=(--ratio="$RATIO")
for f in "${IMAGES[@]}"; do ARGS+=(--images "$f"); done

CRED_BEFORE=$(credit); T0=$(( $(date +%s) - 3 ))
OUTLOG=/tmp/i2i_$$.out; ERRLOG=/tmp/i2i_$$.err
timeout 180 dreamina image2image "${ARGS[@]}" >"$OUTLOG" 2>"$ERRLOG"; CLI_RC=$?

SID=$(python3 -c '
import sys,json,re
t=open(sys.argv[1],encoding="utf-8",errors="replace").read()
for m in re.finditer(r"\{.*?\}", t, re.S):
    try: d=json.loads(m.group(0))
    except Exception: continue
    if d.get("submit_id"): print(d["submit_id"]); break
' "$OUTLOG" 2>/dev/null)
SRC="stdout"
if [ -z "$SID" ]; then
  SID=$(python3 - "$T0" "$PROMPT" <<'PY'
import sqlite3, json, sys, os
since, prompt = int(sys.argv[1]), sys.argv[2]
uri = "file:%s?mode=ro" % os.path.expanduser("~/.dreamina_cli/tasks.db")
db = sqlite3.connect(uri, uri=True)
for sid, req in db.execute("SELECT submit_id, request FROM aigc_task WHERE create_time >= ? ORDER BY create_time DESC", (since,)):
    try: body = json.loads(json.loads(req)["body"])
    except Exception: continue
    if body.get("Prompt") == prompt or body.get("prompt") == prompt:
        print(sid); break
PY
)
  SRC="tasks.db"
fi
[ -n "$SID" ] || { echo "!! 未取得 submit_id（CLI rc=$CLI_RC）" >&2; tail -4 "$ERRLOG" >&2; exit 1; }
echo "submit_id=$SID (via $SRC)"

TMPD=$(mktemp -d); prev=0; DL_OK=0
for i in $(seq 1 60); 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)
  if [ "$n" -gt 0 ]; then
    sz=$(du -sb "$TMPD" | cut -f1)
    if [ "$sz" = "$prev" ] && python3 - "$TMPD" <<'PY'
import sys, os
for f in os.listdir(sys.argv[1]):
    p = os.path.join(sys.argv[1], f)
    if not os.path.isfile(p): continue
    with open(p, "rb") as fh:
        head = fh.read(16); fh.seek(max(0, os.path.getsize(p) - 32)); tail = fh.read()
    if head.startswith(b"\x89PNG"):
        if b"IEND" not in tail: sys.exit(1)
    elif head.startswith(b"\xff\xd8"):
        if not tail.rstrip(b"\x00").endswith(b"\xff\xd9"): sys.exit(1)
    else: sys.exit(1)
sys.exit(0)
PY
    then DL_OK=1; break; fi
    prev="$sz"
  fi
  sleep 5
done
[ "$DL_OK" = "1" ] || { echo "!! 超时未下载到图片 submit_id=$SID" >&2; exit 1; }

FILES=$(find "$TMPD" -type f \( -name '*.png' -o -name '*.jpg' -o -name '*.jpeg' \) | sort)
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=$(credit)
echo "credit: $CRED_BEFORE -> $CRED_AFTER (扣 $((CRED_BEFORE-CRED_AFTER)) CRED_AFTER=${CRED_AFTER})"
echo "submit_id: $SID"