#!/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/AXEL/비주얼/face_sheet.png"
DATASET_DIR = Path.home() / "ai/experiments/axel-dataset-gen/dataset/cropped"
WORKFLOW = "i2i-qwen-image-edit-lightning-4step"  # spark · 1080p 소스 (mac MPS 1080×1440 OOM 위험)
TRIGGER = "axlx"

# 체형 anatomy v3 (2026-04-25 · 남성 lean swimmer-dancer · K-pop lean 미학)
# 규약: 프레임에 보이는 부위만 묘사. 상/하반신 오염 금지 (LoRA 오염 방지)
#   close-up: 체형 묘사 없음 (얼굴만)
#   bust:     어깨·쇄골·상체·가슴
#   thigh-up: 어깨·쇄골·허리·복부·엉덩이까지
#   full-body: 전신 (키·어깨·허리·복부·엉덩이·허벅지·종아리·8.5-9 head)
BODY_ANATOMY = {
    "close-up": "",
    "bust": "Body: lean swimmer-dancer build, broad but softly sculpted shoulders, long horizontal clavicles, naturally lean chest.",
    "thigh-up": "Body: lean swimmer-dancer build, broad shoulders softly tapering to narrow waist, long horizontal clavicles, naturally lean chest, flat tight abdomen, firm glutes from years of dance training.",
    "full-body": "Body: approximately 183cm tall, lean swimmer-dancer build, broad but softly sculpted shoulders tapering to narrow waist, long horizontal clavicles, naturally lean chest, flat tight abdomen, firm glutes, slender long toned legs, 8.5-to-9 head elongated proportion.",
}
BODY_CAPTION_TAIL = {
    "close-up": "",
    "bust": "lean swimmer dancer upper build, broad soft shoulders, long clavicles, lean chest",
    "thigh-up": "lean swimmer dancer build, broad soft shoulders tapering to narrow waist, long clavicles, lean chest, flat abdomen, firm glutes",
    "full-body": "183cm tall, lean swimmer dancer build, broad soft shoulders tapering to narrow waist, long clavicles, lean chest, flat abdomen, firm glutes, slender long legs, 8.5-head proportion",
}
CAPTION_IDENTITY = f"{TRIGGER}, young Korean man"

# 1080p source class (2026-04-26 SKILL §3.2)
SIZE_PORTRAIT = (1080, 1440)  # 3:4 portrait, 1.55MP — close-up/bust/thigh
SIZE_TALL     = (1080, 1440)  # mac mps margin 위해 동일 (1080×1920 OOM 위험)
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

# Orphan recovery map — axel 첫 배치는 빈 상태. 중단 후 사용 시 {job_id: stem_with_suffix} 채움.
RECOVERY: dict[int, str] = {}
_LEGACY_WHISPER_RECOVERY: dict[int, str] = {  # 미사용. axel 에 영향 없음. 폐기 가능.
    # whisper B 세트 잡 ID (보존만)
    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)
# AXEL = mixed-race Korean-Nigerian-American man · 22yo · warm hip-hop R&B · short tousled hair
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. Short tousled black hair. Direct eye contact with the lens.",
     "studio portrait, gray seamless backdrop, cool north window light, short tousled hair, 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. Short hair catching warm highlights.",
     "three-quarter face angle, warm tungsten key light, amber bokeh, short hair warm highlights"),
    ("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. Short hair, defined jawline. Charcoal gray studio wall.",
     "clean side profile, rim light from behind, defined jawline, charcoal gray wall"),
    ("closeup_04_golden_curtain",      "close-up", SIZE_PORTRAIT, EXPR_WARM,
     "Golden hour daylight streaming through a sheer beige linen curtain. Short hair slightly tousled by a breeze. Softly out-of-focus curtain background.",
     "golden hour daylight, sheer linen curtain, wind-tousled short 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. Short hair, faint closed-mouth smile. Near-monochrome gray wall.",
     "contemplative downcast gaze, cool overcast window light, short hair, monochrome gray wall"),
    ("closeup_07_cliff_breeze",        "close-up", SIZE_PORTRAIT, EXPR_WARM,
     "Standing outdoors on a cliff edge, ocean breeze ruffling short hair. Late afternoon sun low and warm. Sea-spray haze in background.",
     "outdoor cliff edge, ocean breeze, ruffled short hair, 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 crew-neck tee. Soft frontal softbox lighting.",
     "white cyclorama studio, white cotton crew neck tee, soft frontal softbox"),
    ("bust_02_evening_denim",          "bust",     SIZE_PORTRAIT, EXPR_WARM,
     "Golden hour outdoor evening. Oversized blue denim trucker jacket layered over a plain white tank. Leaning on a weathered wooden bench. Warm city glow blurred behind.",
     "golden hour outdoor evening, oversized denim trucker jacket, white tank, wooden bench, warm city bokeh"),
    ("bust_03_cafe_knit",              "bust",     SIZE_PORTRAIT, EXPR_WARM,
     "Cafe interior with warm tungsten pendant lamps. Oversized cream chunky knit pullover sweater. Elbows resting on a pale oak table. Steaming coffee cup blurred in the foreground.",
     "cafe interior, warm tungsten pendants, cream chunky knit pullover, oak table, coffee cup foreground blur"),
    ("bust_04_backstage_bomber",       "bust",     SIZE_PORTRAIT, EXPR_COMPOSED,
     "Backstage dressing room. Hollywood mirror bulbs rim-light the silhouette. Cropped black leather varsity bomber with silver snap closures, fine silver chain necklace at the throat.",
     "backstage dressing room, hollywood mirror bulbs, leather varsity bomber, silver snaps, silver chain necklace"),
    ("bust_05_neon_alley_hoodie",      "bust",     SIZE_PORTRAIT, EXPR_COMPOSED,
     "Seoul back alley at night between old brick buildings. Pink and cyan neon signs cast contrasting rim light. Oversized black streetwear hoodie under an open matte black bomber.",
     "seoul back alley night, pink cyan neon rim light, oversized black hoodie, matte black bomber"),
    ("bust_06_rooftop_buttondown",     "bust",     SIZE_PORTRAIT, EXPR_WARM,
     "Rooftop garden in mid-morning. Crisp white linen button-down shirt with collar undone, sleeves rolled to the forearms. A concrete planter out of focus in the foreground. Soft-blurred downtown skyline behind.",
     "rooftop garden morning, white linen button-down, collar undone, rolled sleeves, soft skyline background"),
    ("bust_07_bedroom_henley",         "bust",     SIZE_PORTRAIT, EXPR_WARM,
     "Soft morning sun through linen curtains. Oversized cream waffle henley, top buttons undone. Seated against the headboard of an unmade bed with rumpled white sheets.",
     "morning sun linen curtain, cream waffle henley, top buttons undone, unmade bed, rumpled white sheets"),

    # ── Thigh-up × 6 ──────────────────────────────────
    ("thigh_01_street_loose",          "thigh-up", SIZE_PORTRAIT, EXPR_COMPOSED,
     "Walking along a repaved Seoul side street at dusk. Loose-fit dark blue jeans and a cropped charcoal cardigan over a white tank. Hands tucked into the cardigan pockets.",
     "seoul side street dusk, loose dark blue jeans, cropped charcoal cardigan, white tank, hands in pockets"),
    ("thigh_02_stage_silver_bomber",   "thigh-up", SIZE_PORTRAIT, EXPR_COMPOSED,
     "On a black-floored stage under purple and amber wash lights from above. Cropped silver metallic bomber over a fitted black tank, low-rise black trousers cinched at the waist. Weight shifted onto one leg, mic stand blurred in the foreground.",
     "black stage floor, purple amber wash lights, silver metallic bomber, fitted black tank, low rise black trousers, mic stand foreground"),
    ("thigh_03_window_smartcasual",    "thigh-up", SIZE_PORTRAIT, EXPR_WARM,
     "Seated on a tall stool at a cafe window-bar. Smart-casual layered look — fitted charcoal turtleneck under an unstructured camel blazer, slim olive trousers. Espresso cup on the bar.",
     "cafe window bar stool, charcoal turtleneck, camel blazer, slim olive trousers, espresso cup"),
    ("thigh_04_hotel_silk",            "thigh-up", SIZE_PORTRAIT, EXPR_WARM,
     "Luxury hotel suite in late afternoon. Matching slate-gray silk pajama set with ivory satin piping, top buttons open. Hand resting on a tufted velvet chair back. Gold floor-lamp glow.",
     "luxury hotel suite, slate gray silk pajama set, ivory satin piping, tufted velvet chair, gold floor lamp"),
    ("thigh_05_garden_linen",          "thigh-up", SIZE_PORTRAIT, EXPR_WARM,
     "Manicured rose garden at noon. White linen camp-collar shirt unbuttoned over a plain white tank, tan linen trousers, leather sandals. Sun flare peeking through a wooden trellis.",
     "rose garden noon, white linen camp collar shirt, white tank, tan linen trousers, leather sandals, wooden trellis sun flare"),
    ("thigh_06_record_band_tee",       "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, brown leather belt. One hand pulling a record sleeve from the bin.",
     "vinyl record store, wood paneling, vintage band tee, forest green corduroy trousers, brown leather belt, hand on record sleeve"),

    # ── Full-body × 10 (체형 학습 핵심) ───────────────
    ("fullbody_01_studio_tee_jeans",   "full-body", SIZE_TALL, EXPR_COMPOSED,
     "Plain neutral gray seamless photography studio. Plain 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_overcoat", "full-body", SIZE_TALL, EXPR_COMPOSED,
     "Crossing a wide pedestrian crossing in late autumn. Long charcoal wool overcoat open over a black turtleneck. Black leather Chelsea boots. Caught mid-stride. Blurred traffic behind.",
     "pedestrian crosswalk late autumn, charcoal wool overcoat, black turtleneck, black chelsea boots, mid stride, blurred traffic"),
    ("fullbody_03_dance_studio",       "full-body", SIZE_TALL, EXPR_COMPOSED,
     "Bright mirrored dance studio with a wooden floor. Black ribbed tank and loose black tapered joggers. White low-top sneakers. Mid-warmup stretch with one arm extended overhead. Floor-to-ceiling windows behind.",
     "mirrored dance studio, wooden floor, black ribbed tank, black tapered joggers, white sneakers, warmup stretch, floor to ceiling windows"),
    ("fullbody_04_park_path_smart",    "full-body", SIZE_TALL, EXPR_WARM,
     "Tree-lined park path in early spring. Fitted navy blazer over a cream cotton shirt, slim beige chinos, brown leather loafers. One hand loosely holding a canvas tote bag. Dappled sunlight across the scene.",
     "park path early spring, fitted navy blazer, cream cotton shirt, slim beige chinos, brown leather loafers, canvas tote bag, dappled sunlight"),
    ("fullbody_05_subway_field",       "full-body", SIZE_TALL, EXPR_COMPOSED,
     "Empty subway platform in evening. Oversized olive military field jacket over a black tee and tapered black slacks. Black low-top boots. Platform fluorescents casting clean horizontal shadows. A train blurred behind.",
     "empty subway platform evening, olive field jacket, black tee, tapered black slacks, black low boots, 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 biker jacket over a plain white tee and loose dark jeans, black sneakers. Hands loose at sides. City bokeh in background.",
     "rooftop night seoul skyline, oversized matte leather biker jacket, white tee, loose dark jeans, black sneakers, city bokeh"),
    ("fullbody_07_beach_linen",        "full-body", SIZE_TALL, EXPR_WARM,
     "Walking barefoot along a quiet sandy beach at golden hour. White linen camp-collar shirt unbuttoned, plain white tank underneath, cream linen drawstring trousers rolled at the ankle. Gentle waves and warm low sun behind.",
     "sandy beach golden hour, white linen camp collar shirt, plain white tank, cream linen trousers, rolled hem, barefoot, gentle waves"),
    ("fullbody_08_backstage_stagewear","full-body", SIZE_TALL, EXPR_COMPOSED,
     "Concrete backstage corridor with exposed overhead pipes. Stage performance look — fitted black cropped tank and black wide-leg leather pants, black combat boots, fine silver chain layers. Walking toward camera. Cool fluorescent overhead lighting.",
     "concrete backstage corridor, exposed pipes, fitted black cropped tank, black wide leg leather pants, black combat boots, silver chains, walking forward, fluorescent lighting"),
    ("fullbody_09_gallery_minimalist", "full-body", SIZE_TALL, EXPR_COMPOSED,
     "White-cube art gallery with a large abstract canvas behind. All-black tailored ensemble — fine wool turtleneck and tailored wool trousers, polished black leather oxfords. One hand in the trouser pocket. Track lighting overhead.",
     "white cube art gallery, abstract canvas background, black wool turtleneck, tailored wool trousers, black leather oxfords, hand in pocket, track lighting"),
    ("fullbody_10_elevator_charcoal",  "full-body", SIZE_TALL, EXPR_COMPOSED,
     "Inside a modern mirrored elevator with warm brass trim. Full monochrome charcoal look — heavy cashmere turtleneck, tailored wide-leg trousers, polished black leather derby shoes. Reflection partially visible in the mirrored wall. Warm ceiling downlight.",
     "mirrored elevator brass trim, monochrome charcoal look, cashmere turtleneck, tailored wide leg trousers, black leather derby shoes, 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", "urgent",
        "--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"axlx-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()