#!/usr/bin/env python3
"""
character_turnaround.py — Z-Image 角色三视图/四视图生成
漫剧角色资产：同角色 正面/侧面/背面（或 3/4 视角）展板，单图一张出。

用法:
  python character_turnaround.py --desc "长黑发红发箍白裙少女" --output char.png
  python character_turnaround.py --desc "...（详细角色描述）" --views 4 --seed 888
  # --desc 支持中英文；建议给足细节：发型/发色/发饰/服装/配饰/鞋/体型
"""
import os, sys, json, urllib.request, time, argparse, uuid

COMFY = "http://127.0.0.1:8188"
COMFY_OUT = os.path.expanduser("~/comfy/ComfyUI/output")

VIEWS_3 = ("front view on the left, 90 degree side profile view in the middle, back view on the right")
# 人物四视图标准：正面 / 背面 / 侧面 / 脸部特写（动画·游戏角色资产规范）
VIEWS_4 = ("character sheet, 16:9 horizontal composition, pure white background, flat lighting no shadow, masterpiece, best quality, "
           "left 1/3 of the image is a face extreme close-up portrait with hair-level facial detail looking at viewer, "
           "right 2/3 of the image shows three full-body views arranged in a horizontal row: front view, 90 degree side profile view, back view, "
           "the same character in every panel, identical face hairstyle and outfit across all views, "
           "full body standing pose for the three right-side views")
# 社区布局：面部特写+三视图（16:9 游戏立绘风）
PORTRAIT = ("character design sheet, horizontal layout divided into four panels: "
            "front view, back view, 90 degree side profile view, and a large close-up portrait of the face, "
            "the SAME character in all four panels, identical face hairstyle and outfit, "
            "plain white background, anime style, clean lineart, high quality, detailed")
# 3D展示台布局（爸爸发的开源格式）：上方三视图 + 下方细节特写镜头组
SHOWCASE = ("3D model display, three views of the character (front view, side view, back view), "
            "clean neutral background, below the three main views are close-up detail shots showing "
            "fabric, clothing details, face and accessories, "
            "detail shots of face, collar, fabric texture, accessories, "
            "modern style, 3D render, high quality, masterpiece")
# 场景四视图（动画场景设计标准）：同一空间的 4 种视角，**纯场景无人物**
# ①正面外观 ②背面外观 ③室内视角 ④鸟瞰俯视 —— 强制 empty scene no people
SCENE_VIEWS = [
    ("front", "front exterior view from street level showing the main facade and entrance, empty scene, no characters, no people"),
    ("back", "back exterior view from behind showing the rear side of the building, empty scene, no characters, no people"),
    ("interior", "interior view inside the empty room showing indoor layout, tables chairs lanterns and architectural details, no characters, no people"),
    ("aerial", "top-down orthographic bird eye view of the whole building layout with rooms courtyard trees stone path, floor plan style, no characters, no people"),
]


def submit_wide(desc, seed, prefix, width=2048, height=1024):
    """Z-Image 宽幅全景图（用于场景切段法）。"""
    prompt = (f"wide panoramic establishing shot, continuous sweeping vista of: {desc}, "
              f"single unbroken scene from left to right, consistent architecture and scenery throughout, "
              f"clean background, anime style, game environment design, high quality, detailed")
    wf = {
        "1": {"class_type": "UNETLoader", "inputs": {"unet_name": "z_image_turbo_bf16.safetensors", "weight_dtype": "default"}},
        "2": {"class_type": "CLIPLoader", "inputs": {"clip_name": "qwen_3_4b.safetensors", "type": "lumina2", "device": "default"}},
        "3": {"class_type": "VAELoader", "inputs": {"vae_name": "ae.safetensors"}},
        "4": {"class_type": "CLIPTextEncode", "inputs": {"clip": ["2", 0], "text": prompt}},
        "5": {"class_type": "CLIPTextEncode", "inputs": {"clip": ["2", 0], "text": NEG_PROMPT}},
        "6": {"class_type": "EmptySD3LatentImage", "inputs": {"width": width, "height": height, "batch_size": 1}},
        "7": {"class_type": "ModelSamplingAuraFlow", "inputs": {"model": ["1", 0], "shift": 3.0}},
        "8": {"class_type": "KSampler", "inputs": {"model": ["7", 0], "seed": seed, "steps": 8, "cfg": 1.0,
            "sampler_name": "res_multistep", "scheduler": "simple",
            "positive": ["4", 0], "negative": ["5", 0], "latent_image": ["6", 0], "denoise": 1.0}},
        "9": {"class_type": "VAEDecode", "inputs": {"samples": ["8", 0], "vae": ["3", 0]}},
        "10": {"class_type": "SaveImage", "inputs": {"images": ["9", 0], "filename_prefix": prefix}},
    }
    data = json.dumps({"prompt": wf}).encode()
    req = urllib.request.Request(f"{COMFY}/prompt", data=data, headers={"Content-Type": "application/json"})
    pid = json.loads(urllib.request.urlopen(req, timeout=15).read())["prompt_id"]
    for _ in range(120):
        time.sleep(2)
        h = json.loads(urllib.request.urlopen(f"{COMFY}/history/{pid}", timeout=8).read())
        if pid in h and h[pid].get("outputs"):
            return h[pid]["outputs"]["10"]["images"][0]["filename"]
        if pid in h and h[pid].get("status", {}).get("status_str") == "error":
            for m in h[pid]["status"].get("messages", []):
                if isinstance(m, list) and len(m) > 1 and m[0] == "execution_error":
                    raise RuntimeError(m[1].get("exception_message", "").strip()[:300])
    raise TimeoutError("全景生成超时")


def scene_panorama_views(desc, seed, output_dir, out_name="scene_4view"):
    """全景图切四段法：宽幅全景 → 切 4 段 → 拼成四视图。
    四段天然来自同一张图，一致性 100% 锁定。"""
    from PIL import Image
    import shutil
    prefix = "eed_scene_pan_" + time.strftime("%H%M%S")
    f = submit_wide(desc, seed, prefix, width=2048, height=1024)
    src = os.path.join(COMFY_OUT, f)
    os.makedirs(output_dir, exist_ok=True)
    im = Image.open(src)
    w, h = im.size
    seg_w = w // 4
    views = []
    for i in range(4):
        seg = im.crop((i * seg_w, 0, (i + 1) * seg_w, h))
        seg_path = os.path.join(output_dir, f"{out_name}_view{i+1}.png")
        seg.save(seg_path)
        views.append(seg_path)
    # 拼四视图展示版
    gap = 16
    canvas = Image.new("RGB", (seg_w * 4 + gap * 3, h), "white")
    for i, v in enumerate(views):
        canvas.paste(Image.open(v), (i * (seg_w + gap), 0))
    final = os.path.join(output_dir, f"{out_name}_combo.png")
    canvas.save(final)
    return final, views


def char_four_views_compose(desc, seed, output_dir, out_name="char_4view"):
    """人物四视图（逐张生成 + PIL 拼图，**精确 4 张**）：
    ①正面 ②背面 ③侧面 ④脸部特写。
    逐张提交 Z-Image 标准出图，角色一致性靠相似 seed 锁定。"""
    import shutil
    from PIL import Image
    os.makedirs(output_dir, exist_ok=True)
    # 每张强制 ONLY ONE single character（避免 Z-Image 自由发挥画多个），简化元素避免歧义
    cmds = [
        ("front", f"character design, ONLY ONE single character, {desc}, front view looking at viewer, full body, A-pose, isolated on plain white background, no other figures, anime style, high quality"),
        ("back", f"character design, ONLY ONE single character, {desc}, back view facing away, full body, A-pose, isolated on plain white background, no other figures, anime style, high quality"),
        ("side", f"character design, ONLY ONE single character, {desc}, 90 degree side profile view facing left, full body, A-pose, isolated on plain white background, no other figures, anime style, high quality"),
        ("face", f"character design, ONLY ONE single character, {desc}, face close-up portrait, head and shoulders only, looking at viewer, isolated on plain white background, no other figures, anime style, high quality"),
    ]
    paths = []
    for i, (key, prompt) in enumerate(cmds):
        f = submit(prompt, 1, seed + i * 13, f"eed_char_{key}_{time.strftime('%H%M%S')}")
        p = os.path.join(output_dir, f"{out_name}_{key}.png")
        shutil.copy(os.path.join(COMFY_OUT, f), p)
        paths.append(p)
        print(f"    ✅ {key}: {p}")
    # 2x2 拼图
    ims = [Image.open(p) for p in paths]
    w, h = ims[0].size
    gap = 12
    canvas = Image.new("RGB", (w * 2 + gap, h * 2 + gap), "white")
    canvas.paste(ims[0], (0, 0))
    canvas.paste(ims[1], (w + gap, 0))
    canvas.paste(ims[2], (0, h + gap))
    canvas.paste(ims[3], (w + gap, h + gap))
    combo = os.path.join(output_dir, f"{out_name}_combo.png")
    canvas.save(combo)
    return combo, paths


def scene_four_views(desc, seed, output_dir, out_name="scene_4view"):
    """场景四视图（专业语义）：同一场景的 ①正面外观 ②背面外观 ③室内 ④鸟瞰。
    统一场景描述串 + 相近 seed 逐张生成，保证是"同一个空间"。"""
    import shutil
    from PIL import Image
    os.makedirs(output_dir, exist_ok=True)
    views = []
    base_seed = seed
    for i, (key, view_word) in enumerate(SCENE_VIEWS):
        prompt = f"{desc}, {view_word}, empty scene no characters no people, the same building and location as the other views, consistent architecture details, anime style, game environment concept art, clean style, high quality, detailed"
        print(f"  🎨 视角{i+1}/4 [{key}]: {view_word[:40]}...")
        f = submit(prompt, 1, base_seed + i * 7, f"eed_scene_{key}_{time.strftime('%H%M%S')}")
        v = os.path.join(output_dir, f"{out_name}_{key}.png")
        shutil.copy(os.path.join(COMFY_OUT, f), v)
        views.append(v)
        print(f"    ✅ {v}")
    # 拼 2x2 四视图展示版
    ims = [Image.open(v) for v in views]
    w, h = ims[0].size
    gap = 12
    canvas = Image.new("RGB", (w * 2 + gap, h * 2 + gap), "white")
    pos = [(0, 0), (w + gap, 0), (0, h + gap), (w + gap, h + gap)]
    for im, (x, y) in zip(ims, pos):
        canvas.paste(im, (x, y))
    combo = os.path.join(output_dir, f"{out_name}_combo.png")
    canvas.save(combo)
    return combo, views


NEG_PROMPT = "nsfw, lowres, bad anatomy, text, error, missing fingers, blurry, distorted, watermark, multiple characters, characters, people, persons, human figures, crowd"

def submit(desc, views, seed, prefix, ctype="char", layout="standard"):
    if ctype == "scene":
        prompt = f"environment concept art, {desc}, anime style, high quality"
    elif layout == "showcase":
        prompt = f"{SHOWCASE}. Character and outfit: {desc}"
    elif layout == "portrait":
        prompt = f"{PORTRAIT}. Character: {desc}"
    else:
        view_part = VIEWS_4 if views >= 4 else VIEWS_3
        prompt = (f"character reference sheet, model sheet, {views} views of the SAME character: "
                  f"{view_part}, identical character design across all views, same face hairstyle outfit, "
                  f"plain white background, anime style, clean lineart, "
                  f"character design sheet, high quality, detailed. Character: {desc}")
    wf = {
        "1": {"class_type": "UNETLoader", "inputs": {"unet_name": "z_image_turbo_bf16.safetensors", "weight_dtype": "default"}},
        "2": {"class_type": "CLIPLoader", "inputs": {"clip_name": "qwen_3_4b.safetensors", "type": "lumina2", "device": "default"}},
        "3": {"class_type": "VAELoader", "inputs": {"vae_name": "ae.safetensors"}},
        "4": {"class_type": "CLIPTextEncode", "inputs": {"clip": ["2", 0], "text": prompt}},
        "5": {"class_type": "CLIPTextEncode", "inputs": {"clip": ["2", 0], "text": NEG_PROMPT}},  # 真正的负面词
        "6": {"class_type": "EmptySD3LatentImage", "inputs": {"width": 1024, "height": 1024, "batch_size": 1}},
        "7": {"class_type": "ModelSamplingAuraFlow", "inputs": {"model": ["1", 0], "shift": 3.0}},
        "8": {"class_type": "KSampler", "inputs": {"model": ["7", 0], "seed": seed, "steps": 8, "cfg": 1.0,
            "sampler_name": "res_multistep", "scheduler": "simple",
            "positive": ["4", 0], "negative": ["5", 0], "latent_image": ["6", 0], "denoise": 1.0}},
        "9": {"class_type": "VAEDecode", "inputs": {"samples": ["8", 0], "vae": ["3", 0]}},
        "10": {"class_type": "SaveImage", "inputs": {"images": ["9", 0], "filename_prefix": prefix}},
    }
    data = json.dumps({"prompt": wf}).encode()
    req = urllib.request.Request(f"{COMFY}/prompt", data=data, headers={"Content-Type": "application/json"})
    pid = json.loads(urllib.request.urlopen(req, timeout=15).read())["prompt_id"]
    for _ in range(90):
        time.sleep(2)
        h = json.loads(urllib.request.urlopen(f"{COMFY}/history/{pid}", timeout=8).read())
        if pid in h and h[pid].get("outputs"):
            return h[pid]["outputs"]["10"]["images"][0]["filename"]
        if pid in h and h[pid].get("status", {}).get("status_str") == "error":
            for m in h[pid]["status"].get("messages", []):
                if isinstance(m, list) and len(m) > 1 and m[0] == "execution_error":
                    raise RuntimeError(m[1].get("exception_message", "").strip()[:300])
    raise TimeoutError("Z-Image 生成超时")


def main():
    ap = argparse.ArgumentParser(description="Z-Image 角色三视图/四视图生成")
    ap.add_argument("--desc", required=True, help="角色描述（细节越足越好）")
    ap.add_argument("--views", type=int, default=3, choices=[3, 4], help="视图数：3或4（仅 char 用）")
    ap.add_argument("--type", dest="ctype", default="char", choices=["char", "scene"], help="char=人物三视图 / scene=场景四视图")
    ap.add_argument("--layout", default="standard", choices=["standard", "portrait", "showcase"], help="char布局：standard=标准三栏 / portrait=面部特写+三视图 / showcase=3D展示台(上三视图+下细节特写)")
    ap.add_argument("--method", default="prompt", choices=["prompt", "panorama", "compose"], help="场景: prompt/panorama; 人物四视图: compose=逐张生成+PIL拼图(精确4张)")
    ap.add_argument("--seed", type=int, default=777)
    ap.add_argument("--output", default="", help="输出路径（默认 cang-ying/outputs/character_<time>.png）")
    args = ap.parse_args()

    prefix = ("eed_scene_" if args.ctype == "scene" else "eed_char_") + time.strftime("%H%M%S")
    label = "场景四视图" if args.ctype == "scene" else f"角色{args.views}视图"
    print(f"🎨 Z-Image 生成 {label}: {args.desc[:60]}...")
    import shutil
    if args.ctype == "char" and args.views >= 4 and args.method == "compose":
        out_dir = os.path.expanduser("~/cang-ying/outputs")
        os.makedirs(out_dir, exist_ok=True)
        final, views = char_four_views_compose(args.desc, args.seed, out_dir)
        print(f"✅ 人物四视图(逐张生成+拼图): {final}")
        print(f"   四张图: {', '.join(os.path.basename(v) for v in views)}")
        return
    if args.ctype == "scene":
        out_dir = os.path.expanduser("~/cang-ying/outputs")
        os.makedirs(out_dir, exist_ok=True)
        if args.method == "panorama":
            final, views = scene_panorama_views(args.desc, args.seed, out_dir)
            print(f"✅ 场景四视图(全景切段,备选): {final}")
        else:
            final, views = scene_four_views(args.desc, args.seed, out_dir)
            print(f"✅ 场景四视图(正面/背面/室内/鸟瞰): {final}")
        print(f"   四张视角图: {', '.join(os.path.basename(v) for v in views)}")
        return
    f = submit(args.desc, args.views, args.seed, prefix, args.ctype, args.layout)
    src = os.path.join(COMFY_OUT, f)
    if not args.output:
        out_dir = os.path.expanduser("~/cang-ying/outputs")
        os.makedirs(out_dir, exist_ok=True)
        args.output = os.path.join(out_dir, f"character_{time.strftime('%m%d_%H%M%S')}.png")
    shutil.copy(src, args.output)
    print(f"✅ {label}: {args.output}")


if __name__ == "__main__":
    main()
