import base64
from html import escape

import gradio as gr

from config import (
    BASE_DIR,
    CATEGORIES,
    CITATION_BIBTEX,
    DEFAULT_MOTION,
    HIGHLIGHT_VIDEOS,
    MODELS,
    MODEL_KEYS,
    RESOURCE_LINKS,
    get_motion_metadata,
    get_motions_by_category,
    get_video_path,
    label_to_key,
)

STATIC_DIRS = [str(BASE_DIR / "videos"), str(BASE_DIR / "assets")]
GRADIO_MAJOR = int(gr.__version__.split(".", 1)[0])
gr.set_static_paths(paths=STATIC_DIRS)

with open(BASE_DIR / "style.css", encoding="utf-8") as css_file:
    CUSTOM_CSS = css_file.read()
with open(BASE_DIR / "assets" / "logo.png", "rb") as logo_file:
    LOGO_DATA_URI = f"data:image/png;base64,{base64.b64encode(logo_file.read()).decode('ascii')}"

_STATIC_PREFIX = "/gradio_api/file=" if GRADIO_MAJOR >= 5 else "/file="
TEASER_SRC = f"{_STATIC_PREFIX}{BASE_DIR / 'assets' / 'teaser.gif'}"


def render_resource_links() -> str:
    cards = []
    for item in RESOURCE_LINKS:
        cards.append(
            f"""
            <a class="mm-link-card" href="{escape(item["url"])}" target="_blank" rel="noreferrer">
                <span class="mm-link-title">{escape(item["label"])}</span>
                <span class="mm-link-note">{escape(item["note"])}</span>
            </a>
            """
        )
    return "".join(cards)


def build_hero_html() -> str:
    return f"""
    <section id="top" class="mm-hero">
        <div class="mm-hero-copy">
            <div class="mm-brand">
                <img src="{LOGO_DATA_URI}" alt="MuscleMimic logo">
                <span>MuscleMimic Research Space</span>
            </div>
            <h1>Checkpoint comparison and policy rollouts</h1>
            <p class="mm-hero-lede">
                MuscleMimic lets us train and evaluate many checkpoints and training recipes
                at scale. Every policy shown here is a <strong>generalist</strong> — a single
                checkpoint that learns to track the full KIT kinesis locomotion set jointly, not one
                motion at a time. Here we showcase a few interesting ones: side-by-side
                checkpoint comparisons and policy rollouts across locomotion, manipulation,
                and dynamic motions.
            </p>
            <div class="mm-nav">
                <a href="#explorer" onclick="event.preventDefault();document.getElementById('explorer')?.scrollIntoView({{behavior:'smooth',block:'start'}});">Checkpoint comparison</a>
                <a href="#highlights" onclick="event.preventDefault();document.getElementById('highlights')?.scrollIntoView({{behavior:'smooth',block:'start'}});">Policy rollouts</a>
                <a href="#resources" onclick="event.preventDefault();document.getElementById('resources')?.scrollIntoView({{behavior:'smooth',block:'start'}});">Paper and release links</a>
                <a href="#citation" onclick="event.preventDefault();document.getElementById('citation')?.scrollIntoView({{behavior:'smooth',block:'start'}});">Citation</a>
            </div>
        </div>
        <div class="mm-hero-media">
            <img src="{TEASER_SRC}" alt="Animated teaser showing MuscleMimic policy rollouts">
        </div>
    </section>
    """


def build_explorer_intro_html() -> str:
    return """
    <section id="explorer" class="mm-section">
        <div class="mm-section-head">
            <p class="mm-overline">Checkpoint Comparison</p>
            <h2>Generalist checkpoints on KIT kinesis locomotion</h2>
            <p>
                Each checkpoint is a single <strong>generalist</strong> policy trained to
                track the full set of KIT kinesis locomotion motions from the paper — not a
                motion-specific model. The dropdown lets you probe the same policy on any
                motion from that training set. The colored character is the reference
                motion and the other character is the rollout.
            </p>
        </div>
    </section>
    """


def build_highlights_html() -> str:
    return """
    <section id="highlights" class="mm-section">
        <div class="mm-section-head">
            <p class="mm-overline">Policy Rollouts</p>
            <h2>Policy rollouts across locomotion, manipulation, and dynamic motions</h2>
            <p>
                Selected examples spanning locomotion, object interaction, and more dynamic motions.
                In each video, the blue or purple character is the reference motion.
            </p>
        </div>
    </section>
    """


def build_highlight_meta_html(tag: str, label: str) -> str:
    return f"""
    <div class="mm-video-meta">
        <span class="mm-tag">{escape(tag)}</span>
        <h3>{escape(label)}</h3>
    </div>
    """


def build_resources_html() -> str:
    return f"""
    <section id="resources" class="mm-section">
        <div class="mm-section-head">
            <p class="mm-overline">Resources</p>
            <h2>Paper and release links</h2>
            <p>
                Primary links for the paper, code, released checkpoint, and retargeted dataset.
            </p>
        </div>
        <div class="mm-link-strip">
            {render_resource_links()}
        </div>
    </section>
    """


def build_citation_html() -> str:
    return f"""
    <section id="citation" class="mm-section">
        <div class="mm-section-head">
            <p class="mm-overline">Citation</p>
            <h2>Cite MuscleMimic</h2>
            <p>
                If you use MuscleMimic in your research, please cite the paper:
            </p>
        </div>
        <pre class="mm-bibtex"><code>{escape(CITATION_BIBTEX)}</code></pre>
    </section>
    """


def build_autoplay_js() -> str:
    return """
    (() => {
        const tryPlay = (video) => {
            if (!video) return;
            if (video.readyState === 0) return;
            const playPromise = video.play();
            if (playPromise && typeof playPromise.catch === "function") {
                playPromise.catch(() => {});
            }
        };

        const primeVideo = (video) => {
            if (!video) return;
            video.muted = true;
            video.defaultMuted = true;
            video.autoplay = true;
            video.loop = true;
            video.playsInline = true;
            video.setAttribute("muted", "");
            video.setAttribute("autoplay", "");
            video.setAttribute("loop", "");
            video.setAttribute("playsinline", "");
            video.setAttribute("webkit-playsinline", "");

            if (video.dataset.mmAutoplayBound !== "1") {
                video.dataset.mmAutoplayBound = "1";
                const replay = () => tryPlay(video);
                video.addEventListener("loadedmetadata", replay);
                video.addEventListener("loadeddata", replay);
                video.addEventListener("canplay", replay);
                video.addEventListener("emptied", () => {
                    requestAnimationFrame(replay);
                });
            }

            tryPlay(video);
        };

        const scan = (root = document) => {
            if (root instanceof HTMLVideoElement) {
                primeVideo(root);
                return;
            }
            root.querySelectorAll?.(".gradio-container video").forEach(primeVideo);
        };

        const boot = () => {
            scan();

            const observer = new MutationObserver((mutations) => {
                for (const mutation of mutations) {
                    if (
                        mutation.type === "attributes" &&
                        mutation.target instanceof HTMLVideoElement
                    ) {
                        primeVideo(mutation.target);
                        continue;
                    }
                    mutation.addedNodes.forEach((node) => {
                        if (!(node instanceof HTMLElement)) return;
                        if (node.tagName === "VIDEO") {
                            primeVideo(node);
                        } else {
                            scan(node);
                        }
                    });
                }
            });

            observer.observe(document.body, {
                childList: true,
                subtree: true,
                attributes: true,
                attributeFilter: ["src"],
            });

            document.addEventListener("visibilitychange", () => {
                if (!document.hidden) {
                    scan();
                }
            });

            window.addEventListener("pageshow", () => {
                scan();
            });
        };

        if (document.readyState === "loading") {
            document.addEventListener("DOMContentLoaded", boot, { once: true });
        } else {
            boot();
        }
    })();
    """


def build_motion_summary_html(motion_label: str) -> str:
    if not motion_label:
        motion_label = DEFAULT_MOTION

    motion_meta = get_motion_metadata(motion_label)

    return f"""
    <div class="mm-panel-head">
        <div>
            <p class="mm-overline">Selected Motion</p>
            <h3>{escape(motion_label)}</h3>
            <p>
                Each generalist checkpoint rolled out on this KIT kinesis locomotion trajectory from
                <strong>{escape(motion_meta["category"])}</strong>.
            </p>
        </div>
        <div class="mm-chip-row">
            <span class="mm-chip">Dataset: {escape(motion_meta["dataset"])}</span>
            <span class="mm-chip">Category: {escape(motion_meta["category"])}</span>
            <span class="mm-chip">Reference: blue / purple actor</span>
        </div>
    </div>
    """


def get_motion_assets(motion_label: str):
    if not motion_label:
        motion_label = DEFAULT_MOTION

    motion_key = label_to_key(motion_label)
    return (
        build_motion_summary_html(motion_label),
        *[get_video_path(model_key, motion_key) for model_key in MODEL_KEYS],
    )


def update_explorer(category: str, current_motion: str):
    choices = get_motions_by_category(category)
    next_motion = current_motion if current_motion in choices else choices[0]
    summary_html, *video_paths = get_motion_assets(next_motion)
    return gr.Dropdown(choices=choices, value=next_motion), summary_html, *video_paths


def update_motion_assets(motion_label: str):
    return get_motion_assets(motion_label)


def build_model_card_head(model_key: str) -> str:
    model = MODELS[model_key]
    return f"""
    <div class="mm-compare-head">
        <p class="mm-overline">Checkpoint</p>
        <h3>{escape(model["label"])}</h3>
        <p class="mm-compare-desc">{escape(model["description"])}</p>
    </div>
    """


blocks_kwargs = {"title": "MuscleMimic Research Space"}
if GRADIO_MAJOR >= 6:
    blocks_kwargs["fill_width"] = True
else:
    blocks_kwargs["css"] = CUSTOM_CSS
    blocks_kwargs["js"] = build_autoplay_js()

with gr.Blocks(**blocks_kwargs) as app:
    if GRADIO_MAJOR >= 6:
        gr.HTML(
            value="",
            head=f"<style>{CUSTOM_CSS}</style>",
            js_on_load=build_autoplay_js(),
            visible=False,
        )
    gr.HTML(build_hero_html())
    gr.HTML(build_explorer_intro_html())

    with gr.Row(elem_classes="mm-controls-shell", equal_height=False):
        category_dd = gr.Dropdown(
            choices=CATEGORIES,
            value="All",
            label="Motion Family",
            info="Filter the KIT / KINESIS set by locomotion subtype.",
            elem_classes="mm-select",
            scale=1,
        )
        motion_dd = gr.Dropdown(
            choices=get_motions_by_category("All"),
            value=DEFAULT_MOTION,
            label="Motion Clip",
            info="Choose a KIT kinesis locomotion trajectory for checkpoint comparison.",
            elem_classes="mm-select",
            scale=1,
        )

    _card_class_suffix = {
        "mm-10m-1": "mm-compare-card-10m",
        "mm-base-e": "mm-compare-card-base-e",
        "mm-100M-s1": "mm-compare-card-100m",
        "mm-fullbody-base": "mm-compare-card-fullbody-base",
    }

    initial = get_motion_assets(DEFAULT_MOTION)
    initial_summary_html = initial[0]
    initial_video_paths = list(initial[1:])

    vids = []
    with gr.Column(elem_classes="mm-compare-panel"):
        selected_motion_html = gr.HTML(initial_summary_html)
        with gr.Row(elem_classes="mm-compare-grid", equal_height=False):
            for model_key, video_path in zip(MODEL_KEYS, initial_video_paths):
                suffix = _card_class_suffix.get(model_key)
                card_classes = ["mm-compare-card"]
                if suffix:
                    card_classes.append(suffix)
                with gr.Column(elem_classes=card_classes, min_width=0):
                    gr.HTML(build_model_card_head(model_key))
                    vids.append(
                        gr.Video(
                            value=video_path,
                            autoplay=True,
                            loop=True,
                            show_label=False,
                            container=False,
                            elem_classes="mm-compare-video",
                        )
                    )

    gr.HTML(build_highlights_html())
    with gr.Row(elem_classes="mm-video-grid", equal_height=False):
        for item in HIGHLIGHT_VIDEOS:
            with gr.Column(elem_classes="mm-video-card", scale=1, min_width=0):
                gr.Video(
                    value=str(BASE_DIR / "videos" / "blog" / item["file"]),
                    autoplay=True,
                    loop=True,
                    show_label=False,
                    container=False,
                    elem_classes="mm-video-shell",
                )
                gr.HTML(build_highlight_meta_html(item["tag"], item["label"]))
    gr.HTML(build_resources_html())
    gr.HTML(build_citation_html())

    category_dd.change(
        fn=update_explorer,
        inputs=[category_dd, motion_dd],
        outputs=[motion_dd, selected_motion_html, *vids],
    )
    motion_dd.change(
        fn=update_motion_assets,
        inputs=[motion_dd],
        outputs=[selected_motion_html, *vids],
    )

if __name__ == "__main__":
    launch_kwargs = {"allowed_paths": STATIC_DIRS}
    if GRADIO_MAJOR >= 6:
        launch_kwargs["css"] = CUSTOM_CSS
        launch_kwargs["js"] = build_autoplay_js()
    app.launch(**launch_kwargs)