import os
import tempfile

import gradio as gr
import spaces

from core import HERE, Demo

DEMO = Demo(load_models=os.environ.get("ROBOTRYON_UI_ONLY") != "1")
CLIPS = DEMO.manifest["clips"]
EMB_KEY = {e["name"]: e["key"] for e in DEMO.manifest["embodiments"]}
EMB_NAME = {v: k for k, v in EMB_KEY.items()}
CLIP_LABEL = {c["key"]: f'{c["name"]} ({c["source"]})' for c in CLIPS}
CLIP_KEY = {v: k for k, v in CLIP_LABEL.items()}
CODE = "https://github.com/ryan-seungyong-lee/RoboTryOn"


def _p(rel):
    return str(HERE / rel)


def _sel(clip_label, emb_name):
    return CLIP_KEY[clip_label], EMB_KEY[emb_name]


def stage1_duration(clip_label, emb_name, *_):
    return 20 + 25 * DEMO.episode(*_sel(clip_label, emb_name))["num_keyframes"]   # A100: ~34 s per keyframe


def stage2_duration(clip_label, emb_name, *_):
    return 40 + 60 * DEMO.episode(*_sel(clip_label, emb_name))["num_clips"]       # A100: ~60 s per clip


@spaces.GPU(duration=stage1_duration)
def _stage1(clip_label, emb_name, progress=gr.Progress(track_tqdm=True)):
    clip, emb = _sel(clip_label, emb_name)
    kf_dir, _ = DEMO.translate(clip, emb, tempfile.mkdtemp(prefix="keyframes_"))
    return str(kf_dir / "grid.jpg"), (clip, emb, str(kf_dir))


@spaces.GPU(duration=stage2_duration, size="xlarge")   # LTX-2.5 needs more than a half H200
def _stage2(clip_label, emb_name, kf_dir, progress=gr.Progress(track_tqdm=True)):
    clip, emb = _sel(clip_label, emb_name)
    return str(DEMO.synthesize(clip, emb, kf_dir, tempfile.mkdtemp(prefix="video_")))


def stage1(clip_label, emb_name):
    return _stage1(clip_label, emb_name)


def stage2(clip_label, emb_name, state):
    if not state or tuple(state[:2]) != _sel(clip_label, emb_name):
        raise gr.Error("Run stage 1 (keyframe translation) for this video and embodiment first.")
    return _stage2(clip_label, emb_name, state[2])


def pick_clip(evt: gr.SelectData):
    return CLIP_LABEL[CLIPS[evt.index]["key"]]


def show_clip(clip_label, emb_name, state):
    clip = DEMO.clips[CLIP_KEY[clip_label]]
    keep = state and tuple(state[:2]) == _sel(clip_label, emb_name)
    return (_p(clip["video"]), clip["caption"],
            gr.update(selected_index=[c["key"] for c in CLIPS].index(clip["key"])),
            *((gr.update(), gr.update()) if keep else (None, None)))


def show_embodiment(clip_label, emb_name, state):
    keep = state and tuple(state[:2]) == _sel(clip_label, emb_name)
    return (_p(DEMO.embodiments[EMB_KEY[emb_name]]["reference"]),
            *((gr.update(), gr.update()) if keep else (None, None)))


def load_example(clip_label, emb_name):
    clip, emb = _sel(clip_label, emb_name)
    ex = next(e for e in DEMO.manifest["examples"] if e["clip"] == clip and e["embodiment"] == emb)
    return _p(ex["grid"]), _p(ex["video"]), (clip, emb, _p(ex["keyframes"]))


CSS = """
#clips .thumbnail-item { aspect-ratio: 16 / 9 !important; }
#clips .thumbnail-item img { object-fit: contain !important; }
#clips .caption-label { font-size: 12px; }
"""

first = CLIPS[0]
with gr.Blocks(title="RoboTryOn", css=CSS) as demo:
    gr.Markdown(
        "# RoboTryOn\n"
        "**Synthesizing Dexterous Robot Demonstrations from Human Videos via Virtual Try-On** · "
        f"[Paper](https://arxiv.org/abs/2611.05163) · "
        f"[Project page](https://ryan-seungyong-lee.github.io/RoboTryOn) · [Code]({CODE})\n\n"
        "The input videos are preprocessed in advance (hand reconstruction, retargeting to each "
        "robot hand, arm masks) due to the time limits of this Space. "
        f"To run RoboTryOn on your own videos, see the [code]({CODE}).")
    state = gr.State(None)

    gr.Markdown("### 1. Human video")
    gallery = gr.Gallery([(_p(c["thumb"]), c["name"]) for c in CLIPS], columns=9,
                         rows=1, height=120, object_fit="contain", allow_preview=False,
                         selected_index=0, show_label=False, elem_id="clips")
    with gr.Row():
        human = gr.Video(_p(first["video"]), label="Human video", interactive=False, height=220,
                         scale=1)
        with gr.Column(scale=2):
            clip_dd = gr.Dropdown(list(CLIP_LABEL.values()), value=CLIP_LABEL[first["key"]],
                                  label="Selected video")
            caption = gr.Textbox(first["caption"], label="Task", interactive=False, lines=3)

    gr.Markdown("### 2. Embodiment")
    with gr.Row():
        emb_radio = gr.Radio(list(EMB_KEY), value=next(iter(EMB_KEY)), label="Target embodiment")
        reference = gr.Image(_p(DEMO.manifest["embodiments"][0]["reference"]),
                             label="Reference", interactive=False, height=180)

    gr.Markdown("### 3. Generate")
    with gr.Row():
        b1 = gr.Button("Stage 1: Keyframe translation", variant="primary")
        b2 = gr.Button("Stage 2: Video synthesis (uses the stage-1 keyframes)", variant="primary")
    grid = gr.Image(label="Keyframes: human (top) and robot (bottom)", type="filepath")
    video = gr.Video(label="Robot video")

    gr.Markdown("### Examples (precomputed)")
    gr.Examples([[CLIP_LABEL[e["clip"]], EMB_NAME[e["embodiment"]]] for e in DEMO.manifest["examples"]],
                inputs=[clip_dd, emb_radio], outputs=[grid, video, state], fn=load_example,
                run_on_click=True, cache_examples=False)

    gallery.select(pick_clip, None, clip_dd)
    clip_dd.change(show_clip, [clip_dd, emb_radio, state], [human, caption, gallery, grid, video])
    emb_radio.change(show_embodiment, [clip_dd, emb_radio, state], [reference, grid, video])
    b1.click(stage1, [clip_dd, emb_radio], [grid, state])
    b2.click(stage2, [clip_dd, emb_radio, state], video)

demo.queue().launch()