#!/usr/bin/env python3
"""teaser · 编辑 case 的「外观编辑 + 动作 / 轨迹控制」编排（maliang venv，CPU）。目标：展示「语义管布局、首帧管外观、按键管动作」
三者可以分开控制。产物都在 outputs/teaser/edit/app/（下称 $A）。

子命令
------
plan       底子（BASES：A 档首帧 + vscore 高 + 无硬伤的双向 rollout，见 outputs/teaser/roll/<name>）→ $A/bases.json，
           外观编辑的首帧任务 → $A/f0/jobs.jsonl（edit_app_f0.py 吃；每个底子 7–8 个变体：时间 / 天气 / 季节整图编辑、
           主角换盔甲 / 换英雄、Boss 换造型、地面换材质）。
galleryplan 角色画廊（GALLERY：3 个底子 × 6 种主角、3 个底子 × 6 种 Boss，kind = char 局部编辑）→ $A/f0/gallery.jsonl；
           rolljobs 会把画好的写进 $A/jobs_extra.jsonl。
rolljobs   已编辑好的首帧 → 双向 rollout 任务 $A/jobs_look.jsonl（rollout.py 吃；ids / prompts = 底子原样，只换首帧；
           每个底子另加一条 orig = 原首帧，同一流水线重跑，宫格对照用）。
mirror     镜像 = 动作控制：底子的 ids 与首帧左右镜像（ids[:, :, ::-1]，左右部件 id 互换：12↔13、14↔15、17↔18、19↔20 ⇒ 镜像后的
           角色仍是解剖上一致的人），按键 a / d 互换、prompts 按新按键重写 ⇒「同一个世界、往反方向走」→ $A/inputs/mirror/<b>/；
           另做 text 换动作：ids / 首帧不动，只把 21 条 prompts 换成别的动作文本（全 standing still / 反向 / 一直攻击）→
           $A/inputs/text/<b>/。任务写进 $A/jobs_act.jsonl。
finalize   每条完成的结果目录补齐：first.png（实际首帧）、actions.npy（[20,5] w a s d attack；编辑改了动作就是新的）、
           prompts.json（21 条）、ids.npz（rollout.py 已写）、meta.json 里加 "edit" 键（编辑类型、编辑帧、编辑参数、底子来源）。
compose    $A/showcase/：每个底子「原版 + 各外观变体」宫格（look_<b>.mp4）、全部外观的总宫格、镜像 / 文本对照并排、
           帧拼图 jpg（$A/showcase/sheets/）。同一世界不同游玩的合成在 edit_sameworld.py compose。

usage（$PY = $M/.venv/bin/python + SPEC §0 的 env；在 $M 下运行）：
  $PY scripts/teaser/edit_app.py plan
  CUDA_VISIBLE_DEVICES=6 $M/envs/img/bin/python scripts/teaser/edit_app_f0.py --jobs outputs/teaser/edit/app/f0/jobs.jsonl
  $PY scripts/teaser/edit_app.py rolljobs && $PY scripts/teaser/edit_app.py mirror
  CUDA_VISIBLE_DEVICES=7 nohup $PY scripts/teaser/rollout.py --jobs outputs/teaser/edit/app/jobs_look.jsonl \\
      outputs/teaser/edit/app/jobs_act.jsonl --compile --watch 600 > outputs/teaser/logs/edit/roll_g7.log 2>&1 &
  $PY scripts/teaser/edit_app.py finalize && $PY scripts/teaser/edit_app.py compose
"""
from __future__ import annotations

import argparse
import glob
import json
import os
import shutil
import sys

import numpy as np

M = "/opt/dlami/nvme/zhiyangdeng/maliang"
HERE = os.path.dirname(os.path.abspath(__file__))
for _p in (f"{M}/code/gc25d", HERE):
    if _p not in sys.path:
        sys.path.insert(0, _p)

A = f"{M}/outputs/teaser/edit/app"
ROLL = f"{M}/outputs/teaser/roll"

# ------------------------------------------------------------------ 编辑变体 ----
GLOBAL = {
    "night": "Change only the time of day of picture 1 to a dark moonlit night: cool deep-blue moonlight over the whole "
             "scene, dark soft night shadows, and a few warm orange lights glowing from existing windows, lanterns and "
             "braziers.",
    "dusk": "Change only the lighting of picture 1 to a warm sunset at dusk: low golden-orange sunlight from one side, long "
            "soft shadows and a pink-purple evening glow over the whole scene.",
    "snow": "Turn picture 1 into deep winter: a thick layer of fresh white snow covers all the open ground, the floors and "
            "paths, the roofs, the tops of walls, the rocks and the plants; frost and icicles; cold soft winter light.",
    "autumn": "Turn picture 1 into late autumn: every tree, bush and patch of grass turns red, orange and golden, and fallen "
              "autumn leaves are scattered over the ground and paths; warm soft autumn light.",
    "spring": "Turn picture 1 into fresh spring: the grass becomes lush bright green, the trees and bushes are covered in "
              "pink and white blossoms and fresh green leaves, petals are scattered on the ground; soft bright morning light.",
    "lava": "Turn picture 1 from ice into fire: every ice surface becomes dark volcanic rock and every frozen water surface "
            "becomes glowing molten orange lava; the ice walls become black basalt with glowing cracks; hot red-orange light.",
}
CHAR = {
    "hero_gold": ("player", "Edit picture 1: the hero warrior in the middle now wears ornate shining golden plate armour "
                            "with a flowing red cape and a golden helmet, still holding the same sword."),
    "hero_swap": ("player", "Edit picture 1: replace the hero warrior in the middle with a different hero: a battle mage in "
                            "long flowing crimson robes with a pointed hood and glowing blue magic runes, holding a glowing "
                            "blue sword in the same hand."),
    "boss_swap": ("boss", "Edit picture 1: replace the big monster boss with a different boss: a huge undead skeleton giant "
                          "made of bleached bones, with glowing blue eyes and tattered black cloth, wielding the same kind of "
                          "big weapon."),
}
# 地面换材质：每个底子各自的 (classes, 提示词)
GROUND = {
    "farm": ((2,), "Change only the grey gravel ground and tracks in picture 1 into bright red brick paving."),
    "egypt_sc": ((2,), "Change only the large cracked sandstone floor slabs in picture 1 into polished black-and-white "
                       "checkered marble tiles."),
    "egypt_fl": ((1,), "Change only the golden desert sand in picture 1 into lush green meadow grass."),
    "tibetan": ((2,), "Change only the light grey stone-paved courtyard floor in picture 1 into a soft green grass lawn."),
    "greek": ((2,), "Change only the white marble floor slabs in picture 1 into dark red terracotta floor tiles."),
    "oasis": ((1,), "Change only the golden desert sand in picture 1 into lush green meadow grass."),
    "samurai": ((2,), "Change only the dark grey stone-paved floor in picture 1 into raked white zen garden gravel."),
    "harvest": ((1,), "Change only the trampled brown-yellow grass ground in picture 1 into fresh white snow."),
    "rice": ((1,), "Change only the green terraced rice fields in picture 1 into golden ripe rice fields ready for harvest."),
    "icecave": ((2,), "Change only the smooth translucent blue ice floor in picture 1 into glowing molten lava rock floor."),
}
SEED = {"night": 11, "dusk": 12, "snow": 13, "autumn": 14, "spring": 15, "lava": 16, "hero_gold": 21, "hero_swap": 22,
        "boss_swap": 23, "ground": 31}
STD = ["night", "dusk", "snow", "autumn", "hero_gold", "hero_swap", "boss_swap", "ground"]
BASES = [   # (短名, 底子 rollout 名, 变体)
    ("farm", "ruins__P9C_default_s9945__000096__s4002__farm@photoreal", STD),
    ("egypt_sc", "stonecircle__P9_golden_s7905__001344__s6005__egypt@cinematic", STD),
    ("egypt_fl", "farmland__R10_s9503__004032__s9006__egypt@photoreal", STD),
    ("tibetan", "maze__R10_s3505__003552__s6__tibetan@photoreal", STD),
    ("greek", "colonnade__P9C_default_s10942__001536__s5009__greek@photoreal", [v for v in STD if v != "boss_swap"]),
    ("oasis", "orchard__R10_s6509__001824__s18__oasis@cinematic", [v for v in STD if v != "boss_swap"]),
    ("samurai", "maze__P9_golden_s2907__003648__s11__samurai@photoreal", STD),
    ("harvest", "farmland__R10_s3512__003072__s9012__autumn_harvest@photoreal",
     ["night", "dusk", "snow", "spring", "hero_gold", "hero_swap", "boss_swap", "ground"]),
    ("rice", "patchwork__R10_s3505__003552__s6__rice_terraces@cinematic", STD),
    ("icecave", "arena__P9_golden_s7905__001344__s7000__icecave@photoreal",
     ["night", "dusk", "lava", "hero_gold", "hero_swap", "boss_swap", "ground"]),
]
EDIT_DESC = {"night": "time of day → moonlit night", "dusk": "time of day → sunset", "snow": "weather → snow",
             "autumn": "season → autumn", "spring": "season → spring", "lava": "material → ice to lava",
             "hero_gold": "hero → golden armour", "hero_swap": "hero → crimson battle mage", "boss_swap":
             "boss → skeleton giant", "ground": "ground material", "orig": "original first frame"}
EDIT_EN = {"orig": "Original", "night": "Night", "dusk": "Sunset", "snow": "Snow", "autumn": "Autumn", "spring": "Spring",
           "lava": "Ice to Lava", "hero_gold": "Golden Armour", "hero_swap": "New Hero", "boss_swap": "New Boss",
           "ground": "New Ground"}
SWAP_LR = {12: 13, 13: 12, 14: 15, 15: 14, 17: 18, 18: 17, 19: 20, 20: 19}


def rp(p: str) -> str:
    return p if os.path.isabs(p) else os.path.join(M, p)


def base_info(short: str, name: str) -> dict:
    meta = json.load(open(f"{ROLL}/{name}/meta.json"))
    j = meta["job"]
    md = os.path.dirname(j["ids"])
    vs = f"{M}/outputs/teaser/vscore/roll/{name}.vscore.json"
    v = json.load(open(vs)) if os.path.exists(vs) else {}
    return {"short": short, "name": name, "roll": f"{ROLL}/{name}", "first": j["first"], "ids": j["ids"],
            "prompts": j["prompts"], "actions": f"{md}/actions.npy", "map_dir": md, "scene": j.get("scene"),
            "map": j.get("map"), "case": os.path.basename(md), "vscore": v.get("teaser"), "flags": v.get("flags")}


def load_bases() -> dict:
    return {b["short"]: b for b in json.load(open(f"{A}/bases.json"))}


# ------------------------------------------------------------------ plan ----
def cmd_plan(a) -> None:
    os.makedirs(f"{A}/f0", exist_ok=True)
    bases = [base_info(s, n) | {"variants": vs} for s, n, vs in BASES]
    json.dump(bases, open(f"{A}/bases.json", "w"), indent=1)
    rows = []
    for b in bases:
        for v in b["variants"]:
            j = {"out": f"{A}/f0/{b['short']}/{v}.png", "src": b["first"], "ids": b["ids"], "seed": SEED[v],
                 "base": b["name"], "short": b["short"], "variant": v}
            if v in GLOBAL:
                j.update(kind="global", prompt=GLOBAL[v])
            elif v in CHAR:
                j.update(kind="char", target=CHAR[v][0], prompt=CHAR[v][1])
            elif v == "ground":
                cls, p = GROUND[b["short"]]
                j.update(kind="ground", classes=list(cls), prompt=p)
            rows.append(j)
    with open(f"{A}/f0/jobs.jsonl", "w") as fh:
        for r in rows:
            fh.write(json.dumps(r, ensure_ascii=False) + "\n")
    print(f"{len(bases)} 个底子 → {A}/bases.json；{len(rows)} 个首帧编辑任务 → {A}/f0/jobs.jsonl")


# ------------------------------------------------------------------ rolljobs ----
def roll_row(name: str, first: str, ids: str, prompts, edit: dict, prio: float = 0, out: str = f"{A}/roll") -> dict:
    return {"name": name, "first": first, "ids": ids, "prompts": prompts, "out": out, "seed": 0, "prio": prio,
            "label": edit.get("label"), "scene": edit.get("scene"), "map": edit.get("map"), "edit": edit}


# ------------------------------------------------------------------ 角色画廊（同一张图 × 多种主角 / Boss） ----
HEROES = {
    "h_paladin": "a holy paladin in shining white-and-gold plate armour with a blue tabard",
    "h_ninja": "a black-clad ninja assassin with a dark hood and a face mask",
    "h_ranger": "an elven ranger in green leather armour with a hooded green cloak",
    "h_barbarian": "a muscular barbarian with bare arms, brown fur pelts and a horned helmet",
    "h_crimson": "a knight in dark crimson plate armour with a long black cape",
    "h_samurai": "a samurai in red-and-black lacquered armour with a horned kabuto helmet",
}
BOSSES = {
    "b_golem": "a huge golem made of mossy grey stone blocks with glowing green cracks",
    "b_demon": "a huge fire demon with dark red skin, big horns and glowing orange cracks",
    "b_troll": "a huge ice troll with pale blue skin, white hair and frost on its shoulders",
    "b_minotaur": "a huge brown-furred minotaur with big curved horns and a nose ring",
    "b_orc": "a huge green-skinned orc warlord in spiked black iron armour",
    "b_insect": "a huge armoured insect monster with a glossy purple carapace and many legs",
}
GALLERY = {"greek": "hero", "samurai": "hero", "tibetan": "hero", "rice": "boss", "farm": "boss", "egypt_sc": "boss"}
GAL_EN = {"h_paladin": "Paladin", "h_ninja": "Ninja", "h_ranger": "Elf Ranger", "h_barbarian": "Barbarian",
          "h_crimson": "Crimson Knight", "h_samurai": "Samurai", "b_golem": "Stone Golem", "b_demon": "Fire Demon",
          "b_troll": "Ice Troll", "b_minotaur": "Minotaur", "b_orc": "Orc Warlord", "b_insect": "Insect Monster"}
EDIT_EN.update(GAL_EN)


def gallery_variants(short: str) -> list[str]:
    return list(HEROES if GALLERY.get(short) == "hero" else BOSSES) if short in GALLERY else []


def cmd_galleryplan(a) -> None:
    """角色画廊的首帧编辑任务 → $A/f0/gallery.jsonl（edit_app_f0 吃，kind = char）。"""
    bases = load_bases()
    rows = []
    for s, who in GALLERY.items():
        b = bases[s]
        for i, v in enumerate(gallery_variants(s)):
            if who == "hero":
                p = (f"Edit picture 1: replace the hero warrior in the middle with a different hero: {HEROES[v]}, holding "
                     "the same sword in the same hand.")
            else:
                p = (f"Edit picture 1: replace the big monster boss with a different boss: {BOSSES[v]}, wielding the same "
                     "kind of big weapon.")
            rows.append({"out": f"{A}/f0/{s}/{v}.png", "src": b["first"], "ids": b["ids"], "seed": 41 + i, "base": b["name"],
                         "short": s, "variant": v, "kind": "char", "target": "player" if who == "hero" else "boss",
                         "prompt": p})
    with open(f"{A}/f0/gallery.jsonl", "w") as fh:
        for r in rows:
            fh.write(json.dumps(r, ensure_ascii=False) + "\n")
    print(f"{len(rows)} 个角色画廊首帧任务 → {A}/f0/gallery.jsonl")


def cmd_rolljobs(a) -> None:
    bases = load_bases()
    rows = []
    for b in bases.values():
        common = {"base": b["name"], "base_roll": b["roll"], "map_dir": b["map_dir"], "scene": b["scene"], "map": b["map"]}
        rows.append(roll_row(f"{b['short']}__orig", b["first"], b["ids"], b["prompts"],
                             common | {"type": "appearance", "variant": "orig", "edit_frame": None,
                                       "label": EDIT_EN["orig"], "desc": EDIT_DESC["orig"]}, prio=2))
        for v in b["variants"]:
            png = f"{A}/f0/{b['short']}/{v}.png"
            if not (os.path.exists(png) and os.path.exists(png[:-4] + ".json")):
                continue
            info = json.load(open(png[:-4] + ".json"))
            j = info["job"]
            ed = common | {"type": "appearance", "variant": v, "kind": j["kind"], "edit_frame": 0, "first_src": b["first"],
                           "edit_prompt": info["prompt"], "edit_seed": j["seed"], "editor": "Qwen-Image-2.1-Turbo",
                           "target": j.get("target"), "classes": j.get("classes"), "label": EDIT_EN[v],
                           "desc": EDIT_DESC[v], "struct": info.get("struct")}
            rows.append(roll_row(f"{b['short']}__{v}", png, b["ids"], b["prompts"], ed, prio=1))
    with open(f"{A}/jobs_look.jsonl", "w") as fh:
        for r in rows:
            fh.write(json.dumps(r, ensure_ascii=False) + "\n")
    print(f"{len(rows)} 个外观 rollout 任务 → {A}/jobs_look.jsonl")
    extra = []                                                       # 角色画廊 → jobs_extra.jsonl
    for s in GALLERY:
        b = bases[s]
        common = {"base": b["name"], "base_roll": b["roll"], "map_dir": b["map_dir"], "scene": b["scene"], "map": b["map"]}
        for v in gallery_variants(s):
            png = f"{A}/f0/{s}/{v}.png"
            if not (os.path.exists(png) and os.path.exists(png[:-4] + ".json")):
                continue
            info = json.load(open(png[:-4] + ".json"))
            j = info["job"]
            ed = common | {"type": "appearance", "variant": v, "kind": "char", "gallery": True, "edit_frame": 0,
                           "first_src": b["first"], "edit_prompt": info["prompt"], "edit_seed": j["seed"],
                           "editor": "Qwen-Image-2.1-Turbo", "target": j["target"], "label": GAL_EN[v],
                           "desc": f"{j['target']} -> {GAL_EN[v]}", "struct": info.get("struct")}
            extra.append(roll_row(f"{s}__{v}", png, b["ids"], b["prompts"], ed, prio=0.8))
    with open(f"{A}/jobs_extra.jsonl", "w") as fh:
        for r in extra:
            fh.write(json.dumps(r, ensure_ascii=False) + "\n")
    print(f"{len(extra)} 个角色画廊 rollout 任务 → {A}/jobs_extra.jsonl")


# ------------------------------------------------------------------ mirror / text ----
def mirror_ids(ids: np.ndarray) -> np.ndarray:
    out = ids[:, :, ::-1].copy()
    lut = np.arange(256, dtype=np.uint8)
    for k, v in SWAP_LR.items():
        lut[k] = v
    return lut[out]


def blocks_to_prompts(blocks20: np.ndarray) -> list[str]:
    import text as TXT
    return TXT.cell_prompts(blocks20)


MIRROR_BASES = ["farm", "egypt_sc", "tibetan", "greek", "samurai", "harvest", "rice", "icecave", "oasis", "egypt_fl"]
TEXT_BASES = ["farm", "tibetan", "greek", "samurai", "harvest", "icecave"]
TEXT_VARIANTS = {   # 名字 → blocks20 的变换（[20,5] w a s d attack）
    "still": lambda b: np.zeros_like(b),
    "reverse": lambda b: b[:, [2, 3, 0, 1, 4]],              # w↔s、a↔d（文本说反方向走）
    "attack": lambda b: np.concatenate([b[:, :4], np.ones_like(b[:, 4:])], 1),
}
TEXT_EN = {"still": "Text: standing still", "reverse": "Text: opposite direction", "attack": "Text: always attacking"}


def cmd_mirror(a) -> None:
    import cv2
    bases = load_bases()
    rows = []
    for s in MIRROR_BASES:
        b = bases[s]
        d = f"{A}/inputs/mirror/{s}"
        os.makedirs(d, exist_ok=True)
        with np.load(b["ids"]) as z:
            ids = z["ids"]
        np.savez_compressed(f"{d}/ids.npz", ids=mirror_ids(ids))
        img = cv2.imread(b["first"])
        cv2.imwrite(f"{d}/first.png", img[:, ::-1])
        blocks = np.load(b["actions"])
        mb = blocks[:, [0, 3, 2, 1, 4]].copy()                       # a↔d
        np.save(f"{d}/actions.npy", mb)
        json.dump(blocks_to_prompts(mb), open(f"{d}/prompts.json", "w"), indent=1)
        common = {"base": b["name"], "base_roll": b["roll"], "map_dir": b["map_dir"], "scene": b["scene"], "map": b["map"]}
        rows.append(roll_row(f"{s}__mirror", f"{d}/first.png", f"{d}/ids.npz", f"{d}/prompts.json",
                             common | {"type": "action", "variant": "mirror", "edit_frame": 0,
                                       "desc": "ids + first frame mirrored left↔right (L/R part ids swapped), keys a↔d",
                                       "label": "Mirrored (A/D swapped)", "actions": f"{d}/actions.npy"}, prio=1.5))
        # 镜像对照：镜像语义 + 原首帧不镜像？不做（首帧与语义矛盾）。原版 = <s>__orig（jobs_look 里已有）
    for s in TEXT_BASES:
        b = bases[s]
        blocks = np.load(b["actions"])
        common = {"base": b["name"], "base_roll": b["roll"], "map_dir": b["map_dir"], "scene": b["scene"], "map": b["map"]}
        for tv, fn in TEXT_VARIANTS.items():
            d = f"{A}/inputs/text/{s}/{tv}"
            os.makedirs(d, exist_ok=True)
            nb = fn(blocks).astype(np.uint8)
            np.save(f"{d}/actions.npy", nb)
            json.dump(blocks_to_prompts(nb), open(f"{d}/prompts.json", "w"), indent=1)
            rows.append(roll_row(f"{s}__text_{tv}", b["first"], b["ids"], f"{d}/prompts.json",
                                 common | {"type": "action", "variant": f"text_{tv}", "edit_frame": 0,
                                           "desc": f"only the 21 action prompts replaced ({tv}); ids and first frame unchanged",
                                           "label": TEXT_EN[tv], "actions": f"{d}/actions.npy"}, prio=0.5))
    with open(f"{A}/jobs_act.jsonl", "w") as fh:
        for r in rows:
            fh.write(json.dumps(r, ensure_ascii=False) + "\n")
    print(f"{len(rows)} 个镜像 / 文本 rollout 任务 → {A}/jobs_act.jsonl")


# ------------------------------------------------------------------ finalize ----
def finalize_dir(od: str) -> bool:
    """补齐一条结果目录：first.png / actions.npy / prompts.json / meta.json 的 edit 键。返回是否改动。"""
    mp = f"{od}/meta.json"
    if not os.path.exists(mp):
        return False
    meta = json.load(open(mp))
    j = meta.get("job") or {}                                         # edit_long.py 的结果没有 job（自己写全了文件和 edit 键）
    ed = j.get("edit") or meta.get("edit") or {}
    changed = False
    if not os.path.exists(f"{od}/first.png") and j.get("first"):
        import cv2
        img = cv2.imread(rp(j["first"]))
        cv2.imwrite(f"{od}/first.png", img)
        changed = True
    if not os.path.exists(f"{od}/prompts.json"):
        json.dump(meta["prompts"], open(f"{od}/prompts.json", "w"), indent=1)
        changed = True
    if not os.path.exists(f"{od}/actions.npy"):
        act = ed.get("actions")
        if not act and ed.get("base_roll"):
            act = f"{ed['map_dir']}/actions.npy" if ed.get("map_dir") else None
        if act and os.path.exists(act):
            shutil.copyfile(act, f"{od}/actions.npy")
            changed = True
    if "edit" not in meta:
        meta["edit"] = ed
        meta["files"] = {"video": "gen.mp4", "cond_ids": "ids.npz", "actions": "actions.npy (uint8 [20,5]: w a s d attack, "
                         "0.25 s blocks; cell j = block j-1, cell 0 = cell 1)", "prompts": "prompts.json (21 cell prompts)",
                         "first_frame": "first.png", "side_by_side": "side.mp4", "sheet": "sheet.jpg"}
        tmp = f"{od}/.meta.fin.tmp"
        json.dump(meta, open(tmp, "w"), ensure_ascii=False, indent=1)
        os.replace(tmp, mp)
        changed = True
    return changed


def cmd_finalize(a) -> None:
    """补齐结果目录 + 写总索引 $A/index.jsonl（每条一行：name、dir、type、variant、label、base / group、文件相对路径）。"""
    n = k = 0
    rows = []
    for root in a.roots:
        for od in sorted(glob.glob(f"{rp(root)}/*/")):
            od = od.rstrip("/")
            if os.path.exists(f"{od}/meta.json"):
                n += 1
                k += finalize_dir(od)
                m = json.load(open(f"{od}/meta.json"))
                ed = m.get("edit") or {}
                rows.append({"name": m["name"], "dir": os.path.relpath(od, A), "type": ed.get("type"),
                             "variant": ed.get("variant"), "label": ed.get("label"), "base": ed.get("base"),
                             "group": ed.get("group"), "source_clip": ed.get("source_clip"), "edit_frame": ed.get("edit_frame"),
                             "files": ["gen.mp4", "ids.npz", "actions.npy", "prompts.json", "first.png", "meta.json"]})
    with open(f"{A}/index.jsonl", "w") as fh:
        for r in rows:
            fh.write(json.dumps(r, ensure_ascii=False) + "\n")
    print(f"finalize：{n} 条完成的结果，补齐 {k} 条；索引 → {A}/index.jsonl")


# ------------------------------------------------------------------ compose ----
def _cases(names: list[str], root: str = f"{A}/roll") -> list[str]:
    return [f"{root}/{n}" for n in names if os.path.exists(f"{root}/{n}/meta.json")]


def sheet(paths: list[str], labels: list[str], out: str, frames=(0, 20, 40, 60, 80), tw: int = 312) -> None:
    """帧拼图：每行一条结果（标签 + 帧 0/20/40/60/80），jpg。"""
    import cv2
    th = round(tw * 480 / 832)
    rows = []
    for p, lab in zip(paths, labels):
        z = np.load(f"{p}/frames.npz")
        idx = list(z["idx"])
        tiles = [cv2.resize(z["rgb"][idx.index(f)], (tw, th), interpolation=cv2.INTER_AREA) for f in frames if f in idx]
        r = np.concatenate(tiles, 1).copy()
        for t, f in enumerate(frames):
            cv2.putText(r, f"f{f}", (t * tw + 4, th - 6), cv2.FONT_HERSHEY_SIMPLEX, 0.45, (255, 255, 255), 1, cv2.LINE_AA)
        cv2.rectangle(r, (0, 0), (min(r.shape[1], 12 + 9 * len(lab)), 20), (0, 0, 0), -1)
        cv2.putText(r, lab, (4, 15), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 255, 255), 1, cv2.LINE_AA)
        rows.append(r)
    os.makedirs(os.path.dirname(out), exist_ok=True)
    cv2.imwrite(out, cv2.cvtColor(np.concatenate(rows, 0), cv2.COLOR_RGB2BGR), [cv2.IMWRITE_JPEG_QUALITY, 88])


def cmd_compose(a) -> None:
    import edit_compose as EC
    EC.compose_app(a)


def main(argv=None) -> None:
    ap = argparse.ArgumentParser(description=__doc__.split("\n")[0])
    sp = ap.add_subparsers(dest="cmd", required=True)
    sp.add_parser("plan")
    sp.add_parser("rolljobs")
    sp.add_parser("galleryplan")
    sp.add_parser("mirror")
    f = sp.add_parser("finalize")
    f.add_argument("--roots", nargs="*", default=[f"{A}/roll", f"{A}/sameworld/roll", f"{A}/long/roll"])
    c = sp.add_parser("compose")
    c.add_argument("--only", nargs="*", default=None, help="只出这些部分：look / act / all_look")
    a = ap.parse_args(argv)
    {"plan": cmd_plan, "rolljobs": cmd_rolljobs, "galleryplan": cmd_galleryplan, "mirror": cmd_mirror, "finalize": cmd_finalize,
     "compose": cmd_compose}[a.cmd](a)


if __name__ == "__main__":
    main()