"""Generate the example workflow JSON files (UI format).

Canonical source for everything in example_workflows/ — edit here and re-run,
never hand-edit the JSON. Usage:

    python scripts/build_example_workflows.py                 # rebuild example_workflows/
    python scripts/build_example_workflows.py --demo DIR \
        --profile-a "<selector label>" --profile-b "<label>"  # numbered demo set, prefilled selectors

The demo set adds the fal video workflows (which need the demo-only
fal_i2v_demo.py custom node and a FAL_KEY) and numbers files to tell a story.
"""

import argparse
import json
from pathlib import Path

SENTINEL = "(no profiles found — train one or paste an ID below)"

PACK_ROOT = Path(__file__).resolve().parent.parent

# Widget orders (must match INPUT_TYPES declaration order, required then optional,
# excluding connection-only inputs like IMAGE and PHOTA_PROFILE):
#   PhotaGenerate/PhotaEdit: prompt, num_outputs, aspect_ratio, resolution, quality,
#                            profile_id, api_key, base_model
#   PhotaProfileSelector: profile, manual_profile_id
#   PhotaTrainProfile: image_urls, training_tier, wait_for_ready, tag,
#                      poll_interval_seconds, timeout_minutes, api_key
#   PhotaProfileGallery: tag, api_key | PhotaProfilePicture/Status: profile_id, api_key


class Graph:
    """Tiny builder for ComfyUI UI-format workflow JSON."""

    def __init__(self):
        self.nodes = []
        self.links = []  # [id, from_node, from_slot, to_node, to_slot, type]

    def node(self, ntype, pos, size, widgets, inputs=None, outputs=None, title=None):
        nid = len(self.nodes) + 1
        n = {
            "id": nid, "type": ntype, "pos": pos, "size": size, "flags": {},
            "order": len(self.nodes), "mode": 0, "inputs": inputs or [],
            "outputs": outputs or [], "properties": {"Node name for S&R": ntype},
            "widgets_values": widgets,
        }
        if title:
            n["title"] = title
        self.nodes.append(n)
        return nid

    def link(self, from_node, from_slot, to_node, to_slot, link_type):
        link_id = len(self.links) + 1
        self.links.append([link_id, from_node, from_slot, to_node, to_slot, link_type])
        # Register on both endpoints (the frontend needs the metadata consistent).
        from_outputs = self.nodes[from_node - 1]["outputs"]
        if from_outputs[from_slot]["links"] is None:
            from_outputs[from_slot]["links"] = []
        from_outputs[from_slot]["links"].append(link_id)
        self.nodes[to_node - 1]["inputs"][to_slot]["link"] = link_id
        return link_id

    def to_json(self):
        return {
            "last_node_id": len(self.nodes), "last_link_id": len(self.links),
            "nodes": self.nodes, "links": self.links,
            "groups": [], "config": {}, "extra": {}, "version": 0.4,
        }


def out_image(name="IMAGE"):
    return {"name": name, "type": "IMAGE", "links": None, "slot_index": 0}


def out_slots(*specs):
    return [{"name": n, "type": t, "links": None, "slot_index": i} for i, (n, t) in enumerate(specs)]


def in_image(name="image"):
    return {"name": name, "type": "IMAGE", "link": None}


def in_profile(name):
    return {"name": name, "type": "PHOTA_PROFILE", "link": None}


def in_widget(name, wtype="STRING"):
    return {"name": name, "type": wtype, "link": None, "widget": {"name": name}}


def gen_widgets(prompt, aspect="1:1"):
    return [prompt, 1, aspect, "1K", "auto", "", "", "auto"]


def selector(g, pos, default, title):
    return g.node(
        "PhotaProfileSelector", pos, [380, 130], [default, ""],
        outputs=out_slots(("profile_id", "STRING"), ("profile", "PHOTA_PROFILE")), title=title,
    )


def generate(g, pos, prompt, aspect="1:1", title=None, with_subject=True):
    inputs = [in_profile("subject_1")] if with_subject else []
    return g.node(
        "PhotaGenerate", pos, [420, 360], gen_widgets(prompt, aspect),
        inputs=inputs, outputs=[out_image()], title=title,
    )


def save(g, pos, prefix, title=None):
    return g.node("SaveImage", pos, [320, 280], [prefix], inputs=[in_image("images")], title=title)


def preview(g, pos, title=None):
    return g.node("PreviewImage", pos, [300, 270], [], inputs=[in_image("images")], title=title)


def stitch(g, pos, direction="right"):
    return g.node(
        "ImageStitch", pos, [320, 170], [direction, True, 8, "white"],
        inputs=[in_image("image1"), in_image("image2")], outputs=[out_image()],
    )


# --- The workflows -----------------------------------------------------------


def wf_generate(default):
    g = Graph()
    s = selector(g, [40, 220], default, "Pick a profile (optional)")
    n = generate(g, [480, 200], "A portrait of [[profile]] on a mountain summit at golden hour.", "1:1")
    sv = save(g, [960, 200], "phota_generate")
    g.link(s, 1, n, 0, "PHOTA_PROFILE")
    g.link(n, 0, sv, 0, "IMAGE")
    return g


def wf_edit(default):
    g = Graph()
    li = g.node("LoadImage", [40, 160], [320, 310], ["example.png", "image"],
                outputs=out_slots(("IMAGE", "IMAGE"), ("MASK", "MASK")), title="Photo to edit")
    s = selector(g, [40, 540], default, "Profile (keeps their identity)")
    e = g.node("PhotaEdit", [480, 240], [420, 380],
               ["Make this photo more vibrant.", 1, "auto", "1K", "auto", "", "", "auto"],
               inputs=[in_image(), in_profile("subject_1")], outputs=[out_image()])
    sv = save(g, [960, 240], "phota_edit")
    g.link(li, 0, e, 0, "IMAGE")
    g.link(s, 1, e, 1, "PHOTA_PROFILE")
    g.link(e, 0, sv, 0, "IMAGE")
    return g


def wf_train_then_generate():
    g = Graph()
    t = g.node(
        "PhotaTrainProfile", [40, 200], [460, 400],
        ["\n".join(f"https://example.com/photo{i}.jpg" for i in range(1, 6)), "fast", True, "", 15, 60, ""],
        outputs=out_slots(("profile_id", "STRING"), ("status", "STRING"), ("profile", "PHOTA_PROFILE")),
        title="1. Train: paste 5-10 photo URLs",
    )
    n = generate(g, [560, 200], "A portrait of [[profile]] as an astronaut on the moon, detailed, cinematic.",
                 "1:1", "2. Generate with the new profile")
    sv = save(g, [1040, 200], "phota_train_then_generate")
    g.link(t, 2, n, 0, "PHOTA_PROFILE")
    g.link(n, 0, sv, 0, "IMAGE")
    return g


def wf_profile_gallery(default):
    g = Graph()
    gal = g.node("PhotaProfileGallery", [40, 120], [380, 160], ["", ""],
                 outputs=out_slots(("cover_photos", "IMAGE"), ("listing", "STRING"), ("count", "INT")),
                 title="Everyone on this account")
    pv1 = preview(g, [480, 80], "Cover photos")
    s = selector(g, [40, 480], default, "Inspect one profile")
    pic = g.node("PhotaProfilePicture", [480, 460], [340, 120], ["", ""],
                 inputs=[in_profile("profile")], outputs=[out_image()])
    pv2 = preview(g, [860, 440], "Their cover photo")
    st = g.node("PhotaProfileStatus", [480, 660], [340, 120], ["", ""],
                inputs=[in_profile("profile")],
                outputs=out_slots(("status", "STRING"), ("message", "STRING")))
    g.link(gal, 0, pv1, 0, "IMAGE")
    g.link(s, 1, pic, 0, "PHOTA_PROFILE")
    g.link(pic, 0, pv2, 0, "IMAGE")
    g.link(s, 1, st, 0, "PHOTA_PROFILE")
    return g


def wf_two_subjects(default_a, default_b):
    g = Graph()
    sa = selector(g, [40, 140], default_a, "subject_1 → [[profile1]]")
    sb = selector(g, [40, 360], default_b, "subject_2 → [[profile2]]")
    n = g.node(
        "PhotaGenerate", [480, 200], [420, 380],
        gen_widgets("[[profile1]] photographing [[profile2]] with a vintage camera in a sunflower field.", "4:3"),
        inputs=[in_profile("subject_1"), in_profile("subject_2")], outputs=[out_image()],
        title="Two subjects, distinct roles",
    )
    sv = save(g, [960, 200], "phota_two_subjects")
    g.link(sa, 1, n, 0, "PHOTA_PROFILE")
    g.link(sb, 1, n, 1, "PHOTA_PROFILE")
    g.link(n, 0, sv, 0, "IMAGE")
    return g


def wf_character_sheet(default):
    g = Graph()
    s = selector(g, [40, 420], default, "Pick a profile")
    pic = g.node("PhotaProfilePicture", [480, 40], [340, 120], ["", ""],
                 inputs=[in_profile("profile")], outputs=[out_image()], title="Cover photo")
    prompts = [
        "A professional business headshot of [[profile]], studio lighting, grey backdrop.",
        "[[profile]] walking through Tokyo at night, neon signs, cinematic street photography.",
        "A renaissance oil painting portrait of [[profile]], ornate gold frame, museum quality.",
    ]
    gens = [generate(g, [480, 220 + 390 * i], p, "1:1") for i, p in enumerate(prompts)]
    st1 = stitch(g, [980, 200])
    st2 = stitch(g, [980, 700])
    st3 = stitch(g, [1360, 450], "down")
    sv = save(g, [1740, 450], "phota_character_sheet", "2x2 character sheet")
    g.link(s, 1, pic, 0, "PHOTA_PROFILE")
    for n in gens:
        g.link(s, 1, n, 0, "PHOTA_PROFILE")
    g.link(pic, 0, st1, 0, "IMAGE")
    g.link(gens[0], 0, st1, 1, "IMAGE")
    g.link(gens[1], 0, st2, 0, "IMAGE")
    g.link(gens[2], 0, st2, 1, "IMAGE")
    g.link(st1, 0, st3, 0, "IMAGE")
    g.link(st2, 0, st3, 1, "IMAGE")
    g.link(st3, 0, sv, 0, "IMAGE")
    return g


def wf_ugc_product_shot(default):
    g = Graph()
    li = g.node("LoadImage", [40, 120], [320, 310], ["product.png", "image"],
                outputs=out_slots(("IMAGE", "IMAGE"), ("MASK", "MASK")), title="1. Your product photo")
    s = selector(g, [40, 500], default, "2. Your creator")
    e = g.node(
        "PhotaEdit", [480, 200], [420, 380],
        ["A candid UGC-style selfie photo: [[profile]] holding this exact product up to the camera, smiling, "
         "in a bright modern kitchen, natural morning light, shot on a phone front camera.",
         1, "3:4", "1K", "auto", "", "", "auto"],
        inputs=[in_image(), in_profile("subject_1")], outputs=[out_image()],
        title="3. Product-in-hand shot",
    )
    sv = save(g, [960, 200], "phota_ugc_shot", "Ad-ready still")
    g.link(li, 0, e, 0, "IMAGE")
    g.link(s, 1, e, 1, "PHOTA_PROFILE")
    g.link(e, 0, sv, 0, "IMAGE")
    return g


def wf_video_keyframes(default):
    g = Graph()
    s = selector(g, [40, 420], default, "Pick your star")
    prompts = [
        ("Shot 1: wide", "Cinematic film still, wide establishing shot: [[profile]] walks alone across a foggy stone bridge at dawn, long coat, teal and orange color grade."),
        ("Shot 2: medium", "Cinematic film still, medium shot: [[profile]] pauses mid-bridge and turns to look back, wind in her hair, fog drifting, teal and orange color grade."),
        ("Shot 3: close-up", "Cinematic film still, close-up: [[profile]] breaks into a quiet smile, golden dawn light on her face, shallow depth of field, teal and orange color grade."),
    ]
    gens = [generate(g, [480, 60 + 390 * i], p, "16:9", t) for i, (t, p) in enumerate(prompts)]
    b1 = g.node("ImageBatch", [980, 250], [240, 80], [], inputs=[in_image("image1"), in_image("image2")],
                outputs=[out_image()])
    b2 = g.node("ImageBatch", [1280, 430], [240, 80], [], inputs=[in_image("image1"), in_image("image2")],
                outputs=[out_image()])
    sv = save(g, [1600, 250], "phota_keyframe", "Keyframes (feed to Wan/Kling i2v)")
    wp = g.node("SaveAnimatedWEBP", [1600, 600], [320, 280], ["phota_animatic", 1.0, False, 90, "default"],
                inputs=[in_image("images")], title="Animatic preview")
    for n in gens:
        g.link(s, 1, n, 0, "PHOTA_PROFILE")
    g.link(gens[0], 0, b1, 0, "IMAGE")
    g.link(gens[1], 0, b1, 1, "IMAGE")
    g.link(b1, 0, b2, 0, "IMAGE")
    g.link(gens[2], 0, b2, 1, "IMAGE")
    g.link(b2, 0, sv, 0, "IMAGE")
    g.link(b2, 0, wp, 0, "IMAGE")
    return g


def wf_family_portrait(default_a, default_b, default_c):
    """Three subjects in one shot: two typed sockets + a third id via the STRING link."""
    g = Graph()
    sa = selector(g, [40, 80], default_a, "subject_1")
    sb = selector(g, [40, 300], default_b, "subject_2")
    sc = selector(g, [40, 520], default_c, "third subject (via profile_id)")
    n = g.node(
        "PhotaGenerate", [500, 200], [420, 380],
        gen_widgets(
            "A warm golden-hour family photo of [[profile]] sitting together on a picnic blanket in a park, "
            "candid laughter, film photography style.",
            "4:3",
        ),
        inputs=[in_profile("subject_1"), in_profile("subject_2"), in_widget("profile_id")],
        outputs=[out_image()],
        title="Everyone in one frame",
    )
    sv = save(g, [980, 200], "phota_family_portrait")
    g.link(sa, 1, n, 0, "PHOTA_PROFILE")
    g.link(sb, 1, n, 1, "PHOTA_PROFILE")
    g.link(sc, 0, n, 2, "STRING")
    g.link(n, 0, sv, 0, "IMAGE")
    return g


# --- fal demo workflows (demo set only; need fal_i2v_demo.py + FAL_KEY) -------


def wf_talking_portrait(default):
    g = Graph()
    s = selector(g, [40, 260], default, "1. Pick who speaks")
    n = generate(g, [480, 200],
                 "Cinematic close-up portrait of [[profile]] looking directly at the camera, soft natural light, shallow depth of field.",
                 "16:9", "2. Generate the portrait (Phota)")
    av = g.node("FalTalkingAvatarDemo", [960, 200], [440, 280],
                ["Hi team! I'm a Phota profile. I was trained from just ten photos, and now I can star in any scene you can describe.",
                 "Friendly_Person"],
                inputs=[in_image()],
                outputs=out_slots(("frames", "IMAGE"), ("video_path", "STRING")),
                title="3. Make them speak (TTS + Kling Avatar)")
    wp = g.node("SaveAnimatedWEBP", [1460, 200], [320, 280], ["phota_talking_preview", 24.0, False, 80, "default"],
                inputs=[in_image("images")], title="Preview (mp4 with audio in output/)")
    g.link(s, 1, n, 0, "PHOTA_PROFILE")
    g.link(n, 0, av, 0, "IMAGE")
    g.link(av, 0, wp, 0, "IMAGE")
    return g


def wf_profile_to_video(default):
    g = Graph()
    s = selector(g, [40, 260], default, "1. Pick your star")
    n = generate(g, [480, 200],
                 "Cinematic film still, wide establishing shot: [[profile]] walks alone across a foggy stone bridge at dawn, long coat, teal and orange color grade.",
                 "16:9", "2. Generate the keyframe (Phota)")
    iv = g.node("FalImageToVideoDemo", [960, 200], [440, 260],
                ["The woman walks slowly across the foggy stone bridge toward the camera as dawn light grows; fog drifts; the camera slowly pushes in. Cinematic, smooth motion.",
                 "5"],
                inputs=[in_image()],
                outputs=out_slots(("frames", "IMAGE"), ("video_path", "STRING")),
                title="3. Animate it (Kling via fal)")
    wp = g.node("SaveAnimatedWEBP", [1460, 200], [320, 280], ["phota_video_preview", 24.0, False, 80, "default"],
                inputs=[in_image("images")], title="Preview")
    g.link(s, 1, n, 0, "PHOTA_PROFILE")
    g.link(n, 0, iv, 0, "IMAGE")
    g.link(iv, 0, wp, 0, "IMAGE")
    return g


def wf_two_keyframes_to_video(default):
    g = Graph()
    s = selector(g, [40, 320], default, "1. Pick your subject")
    n = generate(g, [480, 80],
                 "Cinematic close-up portrait of [[profile]], calm neutral expression, looking directly at the camera, soft window light, blurred cafe background.",
                 "16:9", "2. Start frame (Phota Generate)")
    e = g.node("PhotaEdit", [960, 80], [420, 380],
               ["Keep the exact same framing, lighting, and background. She breaks into a joyful laugh, eyes crinkling, head tilted slightly back.",
                1, "auto", "1K", "auto", "", "", "auto"],
               inputs=[in_image(), in_profile("subject_1")], outputs=[out_image()],
               title="3. End frame (Phota Edit: same shot, new expression)")
    fl = g.node("FalFirstLastToVideoDemo", [1460, 80], [440, 280],
                ["A woman's calm expression melts into a joyful laugh, natural facial motion, subtle head movement.", "720p"],
                inputs=[in_image("start_image"), in_image("end_image")],
                outputs=out_slots(("frames", "IMAGE"), ("video_path", "STRING")),
                title="4. Interpolate (Wan FLF2V)")
    wp = g.node("SaveAnimatedWEBP", [1960, 80], [320, 280], ["phota_flf_preview", 16.0, False, 80, "default"],
                inputs=[in_image("images")], title="Video preview")
    g.link(s, 1, n, 0, "PHOTA_PROFILE")
    g.link(n, 0, e, 0, "IMAGE")
    g.link(s, 1, e, 1, "PHOTA_PROFILE")
    g.link(n, 0, fl, 0, "IMAGE")
    g.link(e, 0, fl, 1, "IMAGE")
    g.link(fl, 0, wp, 0, "IMAGE")
    return g


PACK_SET = {
    "generate": lambda a, b: wf_generate(a),
    "edit": lambda a, b: wf_edit(a),
    "train_then_generate": lambda a, b: wf_train_then_generate(),
    "profile_gallery": lambda a, b: wf_profile_gallery(a),
    "two_subjects": wf_two_subjects,
    "character_sheet": lambda a, b: wf_character_sheet(a),
    "ugc_product_shot": lambda a, b: wf_ugc_product_shot(a),
    "video_keyframes": lambda a, b: wf_video_keyframes(a),
}

DEMO_SET = {
    "01_train_then_generate": lambda a, b: wf_train_then_generate(),
    "02_profile_gallery": lambda a, b: wf_profile_gallery(a),
    "03_two_subjects": wf_two_subjects,
    "04_character_sheet": lambda a, b: wf_character_sheet(a),
    "05_ugc_product_shot": lambda a, b: wf_ugc_product_shot(a),
    "06_video_keyframes": lambda a, b: wf_video_keyframes(a),
    "07_talking_portrait_fal": lambda a, b: wf_talking_portrait(a),
    "08_profile_to_video_fal": lambda a, b: wf_profile_to_video(a),
    "09_two_keyframes_to_video_fal": lambda a, b: wf_two_keyframes_to_video(a),
}

# Added to DEMO_SET in main() because it needs the third selector prefill.


def write_set(workflows, dest: Path, profile_a: str, profile_b: str) -> None:
    dest.mkdir(parents=True, exist_ok=True)
    for name, builder in workflows.items():
        (dest / f"{name}.json").write_text(json.dumps(builder(profile_a, profile_b).to_json(), indent=2))
        print(f"wrote {dest / name}.json")


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--demo", type=Path, help="Also write the numbered demo set to this directory")
    parser.add_argument("--profile-a", default=SENTINEL, help="Selector label prefill for subject 1 (demo set)")
    parser.add_argument("--profile-b", default=SENTINEL, help="Selector label prefill for subject 2 (demo set)")
    parser.add_argument("--profile-c", default=SENTINEL, help="Selector label prefill for the family portrait's third subject")
    args = parser.parse_args()

    write_set(PACK_SET, PACK_ROOT / "example_workflows", SENTINEL, SENTINEL)
    if args.demo:
        demo = dict(DEMO_SET)
        demo["10_family_portrait"] = lambda a, b: wf_family_portrait(a, b, args.profile_c)
        write_set(demo, args.demo, args.profile_a, args.profile_b)


if __name__ == "__main__":
    main()