#!/usr/bin/env python3
"""WHISPER Qwen-Image-2512 LoRA 학습용 데이터셋 생성기.

dockaiq i2i-qwen-image-edit-lightning-4step 로 face_sheet 기반 변주 31장 생성 후
musubi-tuner 요구 포맷(image.png + image.txt 쌍) 으로 저장.

구성비 (lora-character-qwen spec):
  close-up 8 / bust 7 / thigh-up 6 / full-body 10 = 31

실행:
  python generate.py                    # 전체 생성 (기존 .png 스킵 · resume)
  python generate.py --dry-run          # 프롬프트만 출력
  python generate.py --only fullbody_06 # 단일 생성
  python generate.py --recover          # 하드코딩 orphan 잡 회수 후 missing 보강
"""
from __future__ import annotations

import argparse
import re
import shutil
import subprocess
import sys
import time
from pathlib import Path

FACE = Path.home() / "Documents/vault/chosn/people/WHISPER/비주얼/face_sheet.png"
DATASET_DIR = Path.home() / "ai/experiments/whisper-dataset-gen/dataset/cropped"
WORKFLOW = "i2i-qwen-image-edit-lightning-4step"  # 2026-04-26 v2 · spark · 1080p uniform
TRIGGER = "whsp"

# 체형 anatomy (2026-04-24 v2 · 4687번 프롬프트 기반 · 여성 기본)
# 규약: 프레임에 보이는 부위만 묘사. 상반신 컷에 하반신 묘사 금지 (LoRA 오염 방지)
#   close-up: 체형 묘사 없음 (얼굴만)
#   bust:     어깨·쇄골·상체만
#   thigh-up: 어깨·쇄골·허리·힙·엉덩이까지
#   full-body: 전신 (키·어깨·허리·힙·엉덩이·허벅지·종아리·8.5-9 head)
# v3 (2026-04-25 · 헬창 편향 제거 · lean dancer build)
BODY_ANATOMY = {
    "close-up": "",
    "bust": "Body: slim toned dancer build, broad but softly sculpted shoulders, long horizontal clavicles.",
    "thigh-up": "Body: lean dancer build, broad shoulders softly wider than hips, long horizontal clavicles, narrow waist, gently curved glutes from years of dance training.",
    "full-body": "Body: approximately 173cm tall, slim lean dancer build, broad but softly sculpted shoulders subtly wider than hips, long horizontal clavicles, narrow natural waist, gently curved glutes, slender long toned legs, 8.5-to-9 head elongated proportion.",
}
BODY_CAPTION_TAIL = {
    "close-up": "",
    "bust": "slim dancer upper build, broad soft shoulders, long clavicles",
    "thigh-up": "lean dancer build, broad soft shoulders, narrow waist, gently curved glutes",
    "full-body": "173cm tall, slim lean dancer build, broad soft shoulders subtly wider than hips, long clavicles, narrow waist, gently curved glutes, slender long legs, 8.5-head proportion",
}
CAPTION_IDENTITY = f"{TRIGGER}, Korean woman"

# v2 (2026-04-26): 1080p uniform. v1 832×1216 은 dataset/cropped_v1_832 에 백업됨.
SIZE_PORTRAIT = (1080, 1440)
SIZE_TALL     = (1080, 1440)
SIZE_SQUARE = (1024, 1024)

POLL_INTERVAL = 5
POLL_TIMEOUT = 7200           # 2h · mac 큐가 깊을 때 A+B 혼합 대기 흡수
SUBMIT_SPACING = 2.0          # sec between submits (sqlite lock 완화)
SUBMIT_LOCK_RETRY = 3
SUBMIT_LOCK_BACKOFF = 5       # sec, exponential

# 2026-04-24 B 세트 (body-aware v2) — 폴러 1800s timeout 로 중도포기 후 수동 recovery 용
RECOVERY: dict[int, str] = {
    4894: "closeup_01_studio_neutral_b",
    4895: "closeup_02_threequarter_tungsten_b",
    4896: "closeup_03_profile_rim_b",
    4897: "closeup_04_golden_curtain_b",
    4898: "closeup_05_candid_laugh_b",
    4899: "closeup_06_overcast_window_b",
    4900: "closeup_07_cliff_breeze_b",
    4901: "closeup_08_skylight_lookup_b",
    4902: "bust_01_white_cyc_tee_b",
    4903: "bust_02_evening_denim_b",
    4904: "bust_03_cafe_cashmere_b",
    4905: "bust_04_backstage_leather_b",
    4906: "bust_05_neon_alley_b",
    4907: "bust_06_rooftop_silk_b",
    4908: "bust_07_bedroom_oxford_b",
    4909: "thigh_01_street_cardigan_b",
    4910: "thigh_02_stage_silver_b",
    4911: "thigh_03_window_bar_slip_b",
    4912: "thigh_04_hotel_pajama_b",
    4913: "thigh_05_rose_garden_b",
    4914: "thigh_06_record_store_b",
    4915: "fullbody_01_studio_tee_jeans_b",
    4916: "fullbody_02_crosswalk_trench_b",
    4917: "fullbody_03_dance_studio_b",
    4918: "fullbody_04_park_path_pink_b",
    4919: "fullbody_05_subway_olive_b",
    4920: "fullbody_06_rooftop_night_b",
    4921: "fullbody_07_beach_linen_b",
    4922: "fullbody_08_backstage_corridor_b",
    4923: "fullbody_09_gallery_column_b",
    4924: "fullbody_10_elevator_monochrome_b",
}


FRAMING = {
    # close-up: 옷 프레임 밖 · 순수 얼굴 identity 학습 (2026-04-24 CEO 결정)
    "close-up": "Extreme close-up of face and jawline only, cropped tightly above the collarbone, NO shoulders, NO clothing visible in frame. Face fills three-quarters of the frame, rule of thirds composition with eyes on the upper horizontal line.",
    "bust": "Bust shot from the mid-chest up, head in the upper third of the frame with gentle headroom, rule of thirds composition.",
    "thigh-up": "Three-quarter frame from mid-thigh up, head in the upper third, rule of thirds composition with generous vertical space.",
    "full-body": "Full body in frame head to toe with feet visible, wide framing, head in the upper third, rule of thirds composition.",
}

POSE_DEFAULT = "Relaxed natural posture, weight shifted onto one leg, shoulders dropped, caught mid-moment — not posed for a photo."

EXPR_WARM = "Gentle closed-mouth smile with a faint eye crinkle, relaxed gaze."
EXPR_COMPOSED = "Composed calm expression, relaxed mouth, soft alert gaze."
EXPR_LAUGH = "Candid unposed laugh caught mid-moment, eyes crinkled, teeth barely visible."


# (stem, composition, size, expression, scene_desc, caption_tail)
ITEMS: list[tuple[str, str, tuple[int, int], str, str, str]] = [
    # ── Close-up × 8 ──────────────────────────────────
    ("closeup_01_studio_neutral",     "close-up",  SIZE_PORTRAIT, EXPR_COMPOSED,
     "Studio face portrait in front of a plain light gray seamless backdrop. Soft cool north-window key light from camera left. Long natural wavy black hair framing the face. Minimal makeup. Direct eye contact with the lens.",
     "studio portrait, gray seamless backdrop, cool north window light, long black wavy hair, minimal makeup, direct eye contact"),
    ("closeup_02_threequarter_tungsten", "close-up", SIZE_PORTRAIT, EXPR_WARM,
     "Three-quarter face angle turned slightly away from camera. Warm tungsten key light from camera-left with amber bokeh drifting behind. Hair tucked behind one ear.",
     "three-quarter face angle, warm tungsten key light, amber bokeh, hair tucked behind ear"),
    ("closeup_03_profile_rim",         "close-up", SIZE_PORTRAIT, EXPR_COMPOSED,
     "Clean side profile. Hard diffused rim light defines the nose bridge and jawline from behind. Hair pulled back into a low ponytail. Charcoal gray studio wall.",
     "clean side profile, rim light from behind, low ponytail, charcoal gray wall"),
    ("closeup_04_golden_curtain",      "close-up", SIZE_PORTRAIT, EXPR_WARM,
     "Golden hour daylight streaming through a sheer beige linen curtain. Hair slightly tousled by a breeze. Softly out-of-focus curtain background.",
     "golden hour daylight, sheer linen curtain, wind-tousled hair, soft warm background"),
    ("closeup_05_candid_laugh",        "close-up", SIZE_PORTRAIT, EXPR_LAUGH,
     "Caught unposed in the middle of laughing. Natural catchlights in the eyes. Shallow depth of field. Leafy outdoor backdrop in afternoon sun.",
     "unposed candid laugh, natural catchlight, shallow depth of field, leafy outdoor backdrop"),
    ("closeup_06_overcast_window",     "close-up", SIZE_PORTRAIT, EXPR_COMPOSED,
     "Contemplative downcast gaze, lashes visible. Cool overcast window light. Long black hair frames the face. Faint closed-mouth smile. Near-monochrome gray wall.",
     "contemplative downcast gaze, cool overcast window light, long hair framing face, monochrome gray wall"),
    ("closeup_07_cliff_breeze",        "close-up", SIZE_PORTRAIT, EXPR_WARM,
     "Standing outdoors on a cliff edge, the ocean breeze blowing hair across the face which she holds back gently. Late afternoon sun low and warm. Sea-spray haze in background.",
     "outdoor cliff edge, ocean breeze, hand holding hair back, late afternoon warm sun, sea spray haze"),
    ("closeup_08_skylight_lookup",     "close-up", SIZE_PORTRAIT, EXPR_WARM,
     "Face tilted slightly upward. Soft skylight bouncing from a high ceiling as diffused top light. Glinting eyes. Deep focus interior gallery behind.",
     "face tilted up, skylight top light, glinting eyes, interior gallery background"),

    # ── Bust × 7 ──────────────────────────────────────
    ("bust_01_white_cyc_tee",          "bust",     SIZE_PORTRAIT, EXPR_COMPOSED,
     "Clean white cyclorama studio. Fitted plain white cotton tee. Long black hair loose over the shoulders. Soft frontal softbox lighting.",
     "white cyclorama studio, white cotton tee, loose long hair, soft frontal softbox"),
    ("bust_02_evening_denim",          "bust",     SIZE_PORTRAIT, EXPR_WARM,
     "Golden hour outdoor evening. Distressed blue denim jacket layered over a charcoal tank. Leaning on a weathered wooden bench. Warm city glow blurred behind.",
     "golden hour outdoor evening, distressed denim jacket, charcoal tank, wooden bench, warm city bokeh"),
    ("bust_03_cafe_cashmere",          "bust",     SIZE_PORTRAIT, EXPR_WARM,
     "Cafe interior with warm tungsten pendant lamps. Oversized cream cashmere turtleneck. Elbows resting on a pale oak table. Steaming coffee cup blurred in the foreground.",
     "cafe interior, warm tungsten pendants, cream cashmere turtleneck, oak table, coffee cup foreground blur"),
    ("bust_04_backstage_leather",      "bust",     SIZE_PORTRAIT, EXPR_COMPOSED,
     "Backstage dressing room. Hollywood mirror bulbs rim-light the silhouette. Cropped black leather bomber with silver snap closures. One hand adjusting a silver earring.",
     "backstage dressing room, hollywood mirror bulbs, cropped leather bomber, silver snaps, adjusting earring"),
    ("bust_05_neon_alley",             "bust",     SIZE_PORTRAIT, EXPR_COMPOSED,
     "Seoul back alley at night between old brick buildings. Pink and cyan neon signs cast contrasting rim light. Oversized matte black leather biker jacket worn over a sheer gray top.",
     "seoul back alley night, pink cyan neon rim light, matte black biker jacket, sheer gray top"),
    ("bust_06_rooftop_silk",           "bust",     SIZE_PORTRAIT, EXPR_WARM,
     "Rooftop garden in mid-morning. Crisp ivory silk blouse with mother-of-pearl buttons, collar open at the collarbone. A concrete planter out of focus in the foreground. Soft-blurred downtown skyline behind.",
     "rooftop garden morning, ivory silk blouse, mother of pearl buttons, concrete planter foreground, soft skyline background"),
    ("bust_07_bedroom_oxford",         "bust",     SIZE_PORTRAIT, EXPR_WARM,
     "Soft morning sun through linen curtains. Oversized pastel blue men's oxford shirt buttoned halfway. Seated cross-legged on an unmade bed with rumpled white sheets.",
     "morning sun linen curtain, pastel blue oxford shirt, unmade bed, rumpled white sheets"),

    # ── Thigh-up × 6 ──────────────────────────────────
    ("thigh_01_street_cardigan",       "thigh-up", SIZE_PORTRAIT, EXPR_COMPOSED,
     "Walking along a repaved Seoul side street at dusk. Faded straight-cut blue jeans and a cropped charcoal cardigan over a white tank. Hands tucked into the cardigan pockets.",
     "seoul side street dusk, faded blue jeans, cropped charcoal cardigan, white tank, hands in pockets"),
    ("thigh_02_stage_silver",          "thigh-up", SIZE_PORTRAIT, EXPR_COMPOSED,
     "On a black-floored stage under purple and amber wash lights from above. Sequined silver mini dress cinched at the waist. Weight shifted onto one leg, mic stand blurred in the foreground.",
     "black stage floor, purple amber wash lights, sequined silver mini dress, cinched waist, mic stand foreground"),
    ("thigh_03_window_bar_slip",       "thigh-up", SIZE_PORTRAIT, EXPR_WARM,
     "Seated on a tall stool at a cafe window-bar. Cappuccino-beige satin slip dress with thin spaghetti straps. Espresso cup on the bar. Long hair gathered to one side.",
     "cafe window bar stool, cappuccino beige slip dress, thin straps, espresso cup, hair gathered to one side"),
    ("thigh_04_hotel_pajama",          "thigh-up", SIZE_PORTRAIT, EXPR_WARM,
     "Luxury hotel suite in late afternoon. Matching bone-colored silk pajama set with ivory satin piping. Hand resting on a tufted velvet chair back. Gold floor-lamp glow.",
     "luxury hotel suite, bone silk pajama set, ivory satin piping, tufted velvet chair, gold floor lamp"),
    ("thigh_05_rose_garden",           "thigh-up", SIZE_PORTRAIT, EXPR_WARM,
     "Manicured rose garden at noon. White broderie anglaise tiered sundress and flat leather sandals. Sun flare peeking through a wooden trellis. Arms relaxed at sides.",
     "rose garden noon, broderie anglaise sundress, leather sandals, wooden trellis sun flare, arms relaxed"),
    ("thigh_06_record_store",          "thigh-up", SIZE_PORTRAIT, EXPR_COMPOSED,
     "Browsing vinyl in a wood-paneled record store. Vintage oversized band tee tucked into high-waisted forest green corduroy trousers. One hand pulling a record sleeve from the bin.",
     "vinyl record store, wood paneling, vintage band tee, forest green corduroy trousers, hand on record sleeve"),

    # ── Full-body × 10 (체형 학습 핵심) ───────────────
    ("fullbody_01_studio_tee_jeans",   "full-body", SIZE_TALL, EXPR_COMPOSED,
     "Plain neutral gray seamless photography studio. White cotton tee and straight-leg indigo jeans. Barefoot. Arms relaxed at the sides. Even softbox lighting head to toe.",
     "neutral gray seamless studio, white tee, straight leg indigo jeans, barefoot, even softbox lighting, full figure visible"),
    ("fullbody_02_crosswalk_trench",   "full-body", SIZE_TALL, EXPR_COMPOSED,
     "Crossing a wide pedestrian crossing in late autumn. Long camel wool trench coat open over a black turtleneck. Black ankle boots. Caught mid-stride. Blurred traffic behind.",
     "pedestrian crosswalk late autumn, camel wool trench coat, black turtleneck, black ankle boots, mid stride, blurred traffic"),
    ("fullbody_03_dance_studio",       "full-body", SIZE_TALL, EXPR_COMPOSED,
     "Bright mirrored dance studio with a wooden floor. Charcoal ribbed tank and high-waisted black bike shorts. White sneakers. Mid-warmup stretch with one arm extended overhead. Floor-to-ceiling windows behind.",
     "mirrored dance studio, wooden floor, charcoal ribbed tank, black bike shorts, white sneakers, warmup stretch, floor to ceiling windows"),
    ("fullbody_04_park_path_pink",     "full-body", SIZE_TALL, EXPR_WARM,
     "Tree-lined park path in early spring. Soft pastel pink midi dress and off-white loafers. One hand loosely holding a canvas tote bag. Dappled sunlight across the scene.",
     "park path early spring, pastel pink midi dress, off white loafers, canvas tote bag, dappled sunlight"),
    ("fullbody_05_subway_olive",       "full-body", SIZE_TALL, EXPR_COMPOSED,
     "Empty subway platform in evening. Oversized olive field jacket over a fitted black dress. Low-heel black booties. Platform fluorescents casting clean horizontal shadows. A train blurred behind.",
     "empty subway platform evening, olive field jacket, fitted black dress, low heel booties, fluorescent horizontal shadows, blurred train"),
    ("fullbody_06_rooftop_night",      "full-body", SIZE_TALL, EXPR_WARM,
     "Rooftop balcony at night with Seoul skyline behind. Oversized matte black leather jacket over a white tee and loose dark jeans. Hands loose at sides. City bokeh in background.",
     "rooftop night seoul skyline, oversized matte leather jacket, white tee, dark loose jeans, city bokeh"),
    ("fullbody_07_beach_linen",        "full-body", SIZE_TALL, EXPR_WARM,
     "Walking barefoot along a quiet sandy beach at golden hour. Flowing ivory linen maxi dress with a side slit. Hair loose in the breeze. Gentle waves and warm low sun behind.",
     "sandy beach golden hour, ivory linen maxi dress, side slit, barefoot, hair in breeze, gentle waves"),
    ("fullbody_08_backstage_corridor", "full-body", SIZE_TALL, EXPR_COMPOSED,
     "Concrete backstage corridor with exposed overhead pipes. Black performance outfit of a fitted cropped top and high-waist wide-leg pants. Walking toward camera. Cool fluorescent overhead lighting.",
     "concrete backstage corridor, exposed pipes, fitted black cropped top, high waist wide leg pants, walking forward, fluorescent lighting"),
    ("fullbody_09_gallery_column",     "full-body", SIZE_TALL, EXPR_COMPOSED,
     "White-cube art gallery with a large abstract canvas behind her. Minimal black sleeveless column dress and pointed black pumps. One hand tucked into the dress side pocket. Track lighting overhead.",
     "white cube art gallery, abstract canvas background, black sleeveless column dress, pointed black pumps, hand in side pocket, track lighting"),
    ("fullbody_10_elevator_monochrome","full-body", SIZE_TALL, EXPR_COMPOSED,
     "Inside a modern mirrored elevator with warm brass trim. Full monochrome charcoal look — cashmere sweater, tailored wide trousers, loafers. Reflection partially visible in the mirrored wall. Warm ceiling downlight.",
     "mirrored elevator brass trim, monochrome charcoal look, cashmere sweater, tailored wide trousers, loafers, partial reflection, warm ceiling downlight"),
]


def build_prompt(composition: str, expr: str, scene: str) -> str:
    framing = FRAMING[composition]
    body = BODY_ANATOMY[composition]
    parts = [
        "Keep the exact face, hair, and identity from the reference image.",
        framing,
    ]
    if body:
        parts.append(body)
    parts += [
        POSE_DEFAULT,
        expr,
        f"Scene: {scene}",
        "Photorealistic, natural color, subtle filmic grain, 35mm lens feel.",
    ]
    return " ".join(parts)


def build_caption(composition: str, tail: str) -> str:
    body_tail = BODY_CAPTION_TAIL[composition]
    segments = [CAPTION_IDENTITY]
    if body_tail:
        segments.append(body_tail)
    segments += [composition, tail]
    return ", ".join(segments)


def submit_one(prompt: str, size: tuple[int, int], label: str) -> int:
    cmd = [
        "dockaiq", "submit-wf",
        "--workflow-id", WORKFLOW,
        "--prompt", prompt,
        "--image", str(FACE),
        "--width", str(size[0]),
        "--height", str(size[1]),
        "--priority", "normal",
        "--label", label,
    ]
    backoff = SUBMIT_LOCK_BACKOFF
    for attempt in range(SUBMIT_LOCK_RETRY):
        r = subprocess.run(cmd, capture_output=True, text=True)
        if r.returncode == 0:
            m = re.search(r"#(\d+)", r.stdout + r.stderr)
            if not m:
                raise RuntimeError(f"no job id in submit output: {r.stdout} {r.stderr}")
            return int(m.group(1))
        if "database is locked" in (r.stderr + r.stdout) and attempt < SUBMIT_LOCK_RETRY - 1:
            time.sleep(backoff)
            backoff *= 2
            continue
        raise RuntimeError(f"submit failed: {r.stderr}\n{r.stdout}")
    raise RuntimeError("submit exhausted retries")


def poll_one(job_id: int, timeout: int = POLL_TIMEOUT) -> str:
    elapsed = 0
    while elapsed < timeout:
        r = subprocess.run(["dockaiq", "show", str(job_id)], capture_output=True, text=True)
        txt = r.stdout
        state_match = re.search(r"^\s*state:\s*(\S+)", txt, re.M)
        state = state_match.group(1) if state_match else ""
        if state in ("done", "completed", "succeeded"):
            out_match = re.search(r"output_paths:\s*\[\"([^\"]+)\"", txt)
            if not out_match:
                raise RuntimeError(f"job {job_id} done but output_paths missing")
            return out_match.group(1)
        if state in ("failed", "error", "canceled"):
            raise RuntimeError(f"job {job_id} {state}")
        time.sleep(POLL_INTERVAL)
        elapsed += POLL_INTERVAL
    raise TimeoutError(f"job {job_id} timeout after {timeout}s")


def caption_for_stem(stem: str) -> str | None:
    for s, comp, _size, _expr, _scene, tail in ITEMS:
        if s == stem:
            return build_caption(comp, tail)
    return None


def copy_job_output(job_id: int, label_stem: str) -> bool:
    """Poll job and copy output to dataset dir as {label_stem}.png/.txt.
    label_stem 은 suffix 포함 가능 (예: 'closeup_01_studio_neutral_b').
    캡션 lookup 은 suffix 제거 후 ITEMS 매칭."""
    # strip _a / _b suffix for caption lookup
    bare = re.sub(r"_[ab]$", "", label_stem)
    caption = caption_for_stem(bare)
    if caption is None:
        print(f"  WARN: stem '{bare}' 이 ITEMS 에 없음 (스킵)", file=sys.stderr)
        return False
    try:
        out_src = poll_one(job_id)
    except Exception as e:
        print(f"  FAIL #{job_id} {label_stem}: {e}", file=sys.stderr)
        return False
    dst_png = DATASET_DIR / f"{label_stem}.png"
    dst_txt = DATASET_DIR / f"{label_stem}.txt"
    # spark 출력 (/home/inhohur/...) 은 mac 에서 직접 접근 불가 → rsync
    if out_src.startswith("/home/inhohur/"):
        rc = subprocess.run(["rsync", "-az", f"spark:{out_src}", str(dst_png)]).returncode
        if rc != 0:
            print(f"  FAIL rsync #{job_id} {label_stem}: rc={rc}", file=sys.stderr)
            return False
    else:
        shutil.copy(out_src, dst_png)
    dst_txt.write_text(caption + "\n")
    print(f"  saved {dst_png.name} (#{job_id})")
    return True


def recover_orphans() -> set[str]:
    """Poll RECOVERY jobs, copy outputs. Return set of covered stems."""
    covered: set[str] = set()
    if not RECOVERY:
        return covered
    print(f"[recover] {len(RECOVERY)} orphan jobs 회수 중 ...")
    for job_id, stem in sorted(RECOVERY.items()):
        if (DATASET_DIR / f"{stem}.png").exists():
            print(f"  skip #{job_id} {stem} (이미 존재)")
            covered.add(stem)
            continue
        if copy_job_output(job_id, stem):
            covered.add(stem)
    return covered


def run(only: str | None, dry_run: bool, recover: bool, suffix: str = "", skip_closeup: bool = False) -> None:
    if not FACE.exists():
        sys.exit(f"FATAL: face_sheet 없음: {FACE}")
    DATASET_DIR.mkdir(parents=True, exist_ok=True)

    items = ITEMS if not only else [it for it in ITEMS if it[0] == only]
    if only and not items:
        sys.exit(f"FATAL: --only '{only}' 매치 없음")
    if skip_closeup:
        items = [it for it in items if it[1] != "close-up"]

    print(f"[dataset-gen] {len(items)}장 타겟 · workflow={WORKFLOW} · trigger={TRIGGER}")

    if dry_run:
        for stem, comp, size, expr, scene, tail in items:
            p = build_prompt(comp, expr, scene)
            c = build_caption(comp, tail)
            print(f"\n=== {stem} ({comp}, {size[0]}x{size[1]}) ===\n  PROMPT: {p}\n  CAPTION: {c}")
        return

    covered: set[str] = set()
    if recover:
        covered = recover_orphans()
        print(f"[recover] covered={len(covered)} stems\n")

    # filter: skip existing png (suffix 포함).
    # suffix 미지정 시 {stem}.png / {stem}_a.png / {stem}_b.png 중 아무 변형이라도 있으면 skip.
    sfx = f"_{suffix}" if suffix else ""
    def _stem_has_any_file(stem: str) -> bool:
        if suffix:
            return (DATASET_DIR / f"{stem}{sfx}.png").exists()
        return any((DATASET_DIR / f"{stem}{s}.png").exists() for s in ["", "_a", "_b"])
    to_submit = [it for it in items if not _stem_has_any_file(it[0])]
    skip = len(items) - len(to_submit)
    if skip:
        print(f"[dataset-gen] {skip}장 이미 존재 — skip resume")

    if not to_submit:
        print(f"[dataset-gen] 모두 완료 — {len(items)} pairs in {DATASET_DIR}")
        return

    # submit + poll (queue processes serially)
    pending: list[tuple[str, int, str, tuple[int, int]]] = []
    for stem, comp, size, expr, scene, tail in to_submit:
        prompt = build_prompt(comp, expr, scene)
        caption = build_caption(comp, tail)
        label_stem = f"{stem}{sfx}"
        job_id = submit_one(prompt, size, f"whsp-ds-{label_stem}")
        print(f"  enqueued #{job_id}  {label_stem}")
        pending.append((label_stem, job_id, caption, size))
        time.sleep(SUBMIT_SPACING)

    print(f"\n[dataset-gen] submitted {len(pending)} jobs · polling sequentially")

    done = 0
    for label_stem, job_id, caption, size in pending:
        if copy_job_output(job_id, label_stem):
            done += 1

    total = done + len(covered) + skip
    print(f"\n[dataset-gen] complete: {done} new + {len(covered)} recovered + {skip} pre-existing = {total}/{len(items)} pairs in {DATASET_DIR}")


def main():
    p = argparse.ArgumentParser()
    p.add_argument("--only", default=None, help="단일 stem 만 생성 (예: fullbody_06_rooftop_night)")
    p.add_argument("--dry-run", action="store_true", help="프롬프트·캡션만 출력")
    p.add_argument("--recover", action="store_true", help="하드코딩 orphan 잡 회수 후 missing 보강")
    p.add_argument("--suffix", default="", help="파일명에 _{suffix} 덧붙임 (예: 'b' → closeup_01_..._b.png)")
    p.add_argument("--skip-close-up", action="store_true", help="close-up 컷 제외 (B 세트에서 close-up 은 이전 런에서 처리됨)")
    args = p.parse_args()
    run(args.only, args.dry_run, args.recover, args.suffix, args.skip_close_up)


if __name__ == "__main__":
    main()
