from html import escape
from urllib.parse import quote

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,
)

gr.set_static_paths(paths=[str(BASE_DIR / "videos"), str(BASE_DIR / "assets")])
FILE_ROUTE_PREFIX = "/gradio_api/file=" if int(gr.__version__.split(".", 1)[0]) >= 6 else "/file="

with open(BASE_DIR / "style.css", encoding="utf-8") as css_file:
    CUSTOM_CSS = css_file.read()


def static_src(*relative_parts: str) -> str:
    return file_src(BASE_DIR.joinpath(*relative_parts))


def file_src(path) -> str:
    return f"{FILE_ROUTE_PREFIX}{quote(str(path), safe='/')}"


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-brand">
            <img src="{static_src("assets", "logo.png")}" 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.
            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">Checkpoint comparison</a>
            <a href="#highlights">Policy rollouts</a>
            <a href="#resources">Paper and release links</a>
            <a href="#citation">Citation</a>
        </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>Checkpoint comparison on KIT locomotion motions</h2>
            <p>
                Compare MyoFullBody checkpoints on KIT locomotion motions.
                The colored character is the reference motion and the other character is the rollout.
            </p>
        </div>
    </section>
    """


def build_highlights_html() -> str:
    cards = []
    for item in HIGHLIGHT_VIDEOS:
        cards.append(
            f"""
            <article class="mm-video-card">
                <div class="mm-video-shell">
                    <video autoplay loop muted playsinline preload="metadata">
                        <source src="{static_src("videos", "blog", item["file"])}" type="video/mp4">
                    </video>
                </div>
                <div class="mm-video-meta">
                    <span class="mm-tag">{escape(item["tag"])}</span>
                    <h3>{escape(item["label"])}</h3>
                </div>
            </article>
            """
        )

    return f"""
    <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>
        <div class="mm-video-grid">
            {''.join(cards)}
        </div>
    </section>
    """


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_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>
                Side-by-side comparison on a KIT 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>
    """


with gr.Blocks(title="MuscleMimic Research Space") as app:
    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,
            min_width=280,
        )
        motion_dd = gr.Dropdown(
            choices=get_motions_by_category("All"),
            value=DEFAULT_MOTION,
            label="Motion Clip",
            info="Choose a KIT locomotion trajectory for checkpoint comparison.",
            elem_classes="mm-select",
            scale=1,
            min_width=280,
        )

    _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):
                    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())
    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__":
    app.launch(css=CUSTOM_CSS)