"""armnet eval — a one-click, no-setup policy eval on real hardware.

Pick a task, paste (or pick) a policy repo id, press Run, and watch the rollouts
on a real SO-101 cell: live video, and per-rollout success auto-scored
by the instrumented BusyBox (no human judging). The run is shared — anyone who
opens this Space follows the same fleet in real time.

The task selector sits above the tabs and drives all three: which dataset the
instructions point at, which policy the Live tab defaults to, and which board
the Leaderboard shows. Jobs route to any free BusyBox cell, so more than one
eval can run at once; the Live tab lets you pick which cell's feed to watch.
"""

from __future__ import annotations

import asyncio

import gradio as gr

import config
import leaderboard
from leaderboard import LEADERBOARD
from runner import STATE, display_cell_id, fleet_availability

# The text timer is a lightweight heartbeat so any viewer notices a run another
# viewer started. Video takes a different path: one long-lived generator stream
# per viewer, not one Gradio request per frame. Five FPS is enough to make the
# robot motion readable without making every public viewer encode and receive
# all 20 policy-rate frames.
_TEXT_REFRESH_SECONDS = 2.0
_VIDEO_STREAM_SECONDS = 0.2
_VIDEO_IDLE_SLEEP_SECONDS = 0.25


def _resolve_policy(task_label: str, preset_label: str, custom: str) -> str:
    custom = (custom or "").strip()
    if custom:
        return custom
    task = config.task_for(task_label)
    return task.presets.get(preset_label, task.reference_policy)


def _session_id(request: gr.Request | None) -> str | None:
    return getattr(request, "session_hash", None) if request is not None else None


def _status_markdown(snap: dict) -> str:
    lines = []
    task = snap["task"]
    policy = snap["policy"]
    if not policy:
        return "_Nothing running. Pick a policy and press Run._"
    lines.append(f"**Policy:** `{policy}`")
    cell = display_cell_id(snap.get("cell_id"))
    if cell:
        lines.append(f"**Cell:** `{cell}`")
    if task is not None:
        lines.append(f"**Task:** {task.label}")
    lines.append(f"**Status:** {snap['status']}")
    if snap["job_id"]:
        lines.append(f"**Job:** `{snap['job_id']}`")
    lines.append(
        f"**Progress:** {snap['n_scored']}/{config.N_EPISODES} rollouts scored"
        f" · {snap['n_success']} succeeded"
    )
    if snap["error"]:
        lines.append(f"\n⚠️ **Error:** {snap['error']}")
    return "  \n".join(lines)


def _rollout_markdown(snap: dict) -> str:
    """Render the scored rollouts as a Markdown table.

    Markdown (unlike gr.Dataframe here) re-renders reliably on every Timer tick,
    so the table grows live as each rollout is scored.
    """
    rows = snap["rollouts"]
    if not rows:
        return "### Rollouts\n\n_No rollouts scored yet._"
    lines = ["### Rollouts", "", "| Rollout | Result | Scored by |", "| --- | --- | --- |"]
    for r in rows:
        result = "✅ success" if r["success"] else "❌ fail"
        lines.append(f"| {r['episode']} | {result} | {r['scored_by']} |")
    return "\n".join(lines)


def _summary_markdown(snap: dict) -> str:
    summary = snap["summary"]
    if not summary:
        if snap["n_scored"]:
            rate = 100.0 * snap["n_success"] / snap["n_scored"]
            return f"### Live: {snap['n_success']}/{snap['n_scored']} succeeded ({rate:.0f}%)"
        return ""
    agg = summary.get("aggregated", {}) if isinstance(summary, dict) else {}
    pc = agg.get("pc_success")
    parts = []
    if pc is not None:
        parts.append(f"### Final: {pc:.0f}% success over {agg.get('n_episodes', '?')} rollouts")
    url = summary.get("eval_dataset_url") if isinstance(summary, dict) else None
    if url:
        parts.append(f"[Recorded eval dataset →]({url})")
    return "  \n".join(parts)


def _refresh_text(task_label: str, selected_feed: str | None, request: gr.Request):
    """Lightweight heartbeat for status shared between every viewer."""
    task = config.task_for(task_label)
    session_id = _session_id(request)
    if selected_feed:
        STATE.set_view(session_id, selected_feed)
    snap = STATE.snapshot(session_id=session_id, task_slug=task.slug)
    return (
        _status_markdown(snap),
        _rollout_markdown(snap),
        _summary_markdown(snap),
        "\n".join(snap["events"][-30:]),
        _run_button_update(snap, task),
        _cell_light_html(snap, task),
        _queue_markdown(snap),
        _feed_selector_update(snap),
        _feed_caption(snap),
    )


async def _stream_video(request: gr.Request):
    """Yield newest frames over one Gradio stream, capped at five FPS.

    ``STATE`` itself is fed continuously by the orchestrator websocket. Reading
    only its newest version preserves the old drop-to-stay-live semantics: a
    slow browser skips stale frames instead of building latency.

    The generator deliberately stays alive while the page is open, including
    between evals. That lets a viewer who did not press Run see a shared run
    started by someone else without bringing back a polling request per frame.
    Gradio cancels it when that browser disconnects. The selected cell is read
    from this session's view on every tick so switching feeds does not restart
    the stream.
    """
    session_id = _session_id(request)
    last_version = -1
    last_cell = None
    while True:
        snap = STATE.snapshot(session_id=session_id)
        cell = snap.get("selected_cell")
        version = snap["frame_version"]
        if cell != last_cell:
            last_cell = cell
            last_version = -1
            if snap["latest_frame"] is None:
                yield None
                await asyncio.sleep(_VIDEO_IDLE_SLEEP_SECONDS)
                continue
        if version == last_version or snap["latest_frame"] is None:
            # An async sleep is cancellable when the browser disconnects. The
            # old synchronous generator blocked one of Gradio's 40 worker
            # threads inside next() whenever the Space was idle.
            await asyncio.sleep(_VIDEO_IDLE_SLEEP_SECONDS)
            continue
        last_version = version
        # decode_frame() already resized this once process-wide. Do not make
        # every connected viewer repeat the same PIL work.
        yield snap["latest_frame"]
        await asyncio.sleep(_VIDEO_STREAM_SECONDS)


def _load_leaderboard(task_label: str, sort_choice: str, force: bool = False):
    by = "rate" if "Average" in (sort_choice or "") else "lower"
    return LEADERBOARD.rows(config.task_for(task_label).slug, by=by, force=force)


def _on_run(task_label: str, preset_label: str, custom: str, request: gr.Request):
    task = config.task_for(task_label)
    ok, message = STATE.start_run(_resolve_policy(task_label, preset_label, custom), task)
    snap = STATE.snapshot(session_id=_session_id(request), task_slug=task.slug)
    banner = ("✅ " if ok else "⛔ ") + message
    return (
        banner,
        _status_markdown(snap),
        _run_button_update(snap, task),
        _feed_selector_update(snap),
        _feed_caption(snap),
        _cell_light_html(snap, task),
    )


def _on_feed_change(selected_feed: str, task_label: str, request: gr.Request):
    task = config.task_for(task_label)
    STATE.set_view(_session_id(request), selected_feed)
    snap = STATE.snapshot(session_id=_session_id(request), task_slug=task.slug)
    return (
        _status_markdown(snap),
        _rollout_markdown(snap),
        _summary_markdown(snap),
        "\n".join(snap["events"][-30:]),
        _feed_caption(snap),
    )


def _on_task_change(task_label: str, sort_choice: str, request: gr.Request):
    """Re-point every tab at the newly selected task."""
    task = config.task_for(task_label)
    snap = STATE.snapshot(session_id=_session_id(request), task_slug=task.slug)
    return (
        _instructions_markdown(task),
        gr.update(choices=list(task.presets), value=task.reference_label),
        gr.update(placeholder=task.reference_policy, value=""),
        _cell_light_html(snap, task),
        _run_button_update(snap, task),
        _load_leaderboard(task_label, sort_choice),
        _task_blurb(task),
        _feed_selector_update(snap),
        _feed_caption(snap),
    )


def _task_blurb(task: config.Task) -> str:
    return f"**{task.label}.** {task.blurb}"


def _instructions_markdown(task: config.Task) -> str:
    return (
        "## How to use this Space\n\n"
        f"1. **Train a LeRobot-compatible policy** on either the task-specific "
        f"[`{task.dataset}`](https://huggingface.co/datasets/{task.dataset}) dataset "
        f"or the language-conditioned "
        f"[`{config.MULTITASK_DATASET}`]"
        f"(https://huggingface.co/datasets/{config.MULTITASK_DATASET}) dataset. "
        "Make sure the trained policy is **uploaded to Hugging Face and publicly "
        "available**.\n"
        "2. **Copy your repo ID** into the text box on the **Live** tab and submit it. "
        "Armnet will automatically run and score the policy on a real SO-101 cell. "
        "If every cell is busy, the job waits in the shared queue. If more than one "
        "eval is running, pick which cell's live feed to watch.\n\n"
        "You can run a policy more than once. The more rollouts a policy accumulates, the "
        "closer its **lower confidence interval** matches its **average success rate** — "
        "see the **Leaderboard** tab.\n\n"
        f"Every task has its own dataset and its own board, so pick the one you want at "
        "the top of the page before you start."
    )


# Gradio's Dataframe markdown strips <img> sizing attributes, so force a small,
# inline avatar next to the trainer name via CSS. Hugging Face shows these as
# squares, not circles. Also styles the round, blinking cell-availability light.
_CSS = """
#leaderboard-table img {
  width: 20px !important;
  height: 20px !important;
  min-width: 20px !important;
  border-radius: 2px;
  display: inline-block;
  vertical-align: middle;
  margin: 0 4px 0 0 !important;
  object-fit: cover;
}
#leaderboard-table .lb-verified {
  color: #16a34a;
  font-weight: 600;
  white-space: nowrap;
  cursor: help;
}
.cell-status {
  display: flex;
  align-items: center;
  gap: 8px;
  font-weight: 600;
  margin: 4px 0 2px 0;
}
.cell-dot {
  width: 14px;
  height: 14px;
  border-radius: 50%;
  display: inline-block;
  flex: 0 0 auto;
}
.cell-dot--green { background: #22c55e; box-shadow: 0 0 7px 1px #22c55e; animation: cell-blink 1.4s ease-in-out infinite; }
.cell-dot--yellow { background: #eab308; box-shadow: 0 0 7px 1px #eab308; animation: cell-blink 1.0s ease-in-out infinite; }
.cell-dot--red { background: #ef4444; box-shadow: 0 0 7px 1px #ef4444; animation: cell-blink 0.8s ease-in-out infinite; }
.cell-dot--grey { background: #9ca3af; }
@keyframes cell-blink { 0%, 100% { opacity: 1; } 50% { opacity: 0.25; } }
"""


def _cell_light_html(snap: dict, task: config.Task) -> str:
    color = snap.get("availability_color")
    caption = snap.get("availability_caption")
    if not color or not caption:
        color, caption = fleet_availability(
            snap.get("cell_statuses") or {},
            aggregate=snap.get("cell_status", {}).get(task.slug, "offline"),
        )
    return _light_html(color, caption)


def _light_html(color: str, caption: str) -> str:
    return (
        '<div class="cell-status">'
        f'<span class="cell-dot cell-dot--{color}"></span>'
        f"<span>{caption}</span>"
        "</div>"
    )


def _feed_caption(snap: dict) -> str:
    policy = snap.get("policy")
    if not policy:
        return ""
    cell = display_cell_id(snap.get("cell_id"))
    task = snap.get("task")
    parts = [f"**Watching:** `{policy}`"]
    if cell:
        parts.append(f"on `{cell}`")
    if task is not None:
        parts.append(f"· {task.label}")
    return " ".join(parts)


def _feed_selector_update(snap: dict):
    feeds = snap.get("live_feeds") or []
    # Offer a choice whenever two evals can be watched — including a run that
    # just finished while another cell is still going.
    shown = feeds if len(feeds) > 1 else []
    choices = []
    for feed in shown:
        cell = display_cell_id(feed["cell_id"]) or "routing"
        suffix = "" if feed["active"] else f" ({feed['status']})"
        choices.append((f"{feed['policy']} · {cell}{suffix}", feed["cell_id"]))
    selected = snap.get("selected_cell")
    values = [choice[1] for choice in choices]
    value = selected if selected in values else (values[0] if values else None)
    return gr.update(choices=choices, value=value, visible=len(choices) > 1)


def _queue_mode(snap: dict, task: config.Task) -> bool:
    statuses = snap.get("cell_statuses") or {}
    if statuses:
        return not any(
            status == "available" for status in statuses.values()
        ) and any(
            status in {"occupied", "maintenance"} for status in statuses.values()
        )
    return snap.get("cell_status", {}).get(task.slug, "offline") == "occupied"


def _can_submit(snap: dict, task: config.Task) -> bool:
    """Allow immediate runs, or queueing while every matching cell is busy."""
    if snap["queued_eval_total"] >= config.MAX_QUEUED_JOBS:
        return False
    statuses = snap.get("cell_statuses") or {}
    if statuses:
        return any(
            status in {"available", "occupied", "maintenance"}
            for status in statuses.values()
        )
    status = snap.get("cell_status", {}).get(task.slug, "offline")
    return status in {"available", "occupied"}


def _run_button_update(snap: dict, task: config.Task):
    queue_mode = _queue_mode(snap, task)
    return gr.update(
        value=(
            f"Queue {config.N_EPISODES} rollouts"
            if queue_mode
            else f"Run {config.N_EPISODES} rollouts"
        ),
        variant="secondary" if queue_mode else "primary",
        interactive=_can_submit(snap, task),
    )


def _queue_markdown(snap: dict) -> str:
    total = snap["queue_total"]
    eval_total = snap["queued_eval_total"]
    lines = [
        f"### Shared queue ({total} waiting)",
        "",
        f"_Armnet eval submissions: {eval_total}/{config.MAX_QUEUED_JOBS}._",
        "",
    ]
    if not snap["queue_items"]:
        lines.append("_No jobs waiting._")
    else:
        lines.extend(
            f"{index}. `{label}`"
            for index, label in enumerate(snap["queue_items"], start=1)
        )
    if eval_total >= config.MAX_QUEUED_JOBS:
        lines.extend(["", "_Armnet eval queue full — submissions are temporarily disabled._"])
    return "\n".join(lines)

_DEFAULT_TASK = config.TASKS[config.DEFAULT_TASK_LABEL]

with gr.Blocks(title="armnet eval — evaluate a policy on real robots") as demo:
    # Inject CSS via a <style> element (Gradio 6 moved the Blocks `css=` arg to
    # launch(), which the Space controls, so this is the portable way).
    gr.HTML(f"<style>{_CSS}</style>")
    gr.Markdown(
        "# armnet eval — evaluate a policy on real robots\n"
        "[armnet.dev](https://armnet.dev/)\n\n"
        "Pick a task, pick or paste a policy, press **Run**, and watch **"
        f"{config.N_EPISODES} rollouts** on a real SO-101 cell. Success is "
        "auto-scored by the instrumented BusyBox — no human judging. Anyone here "
        "can watch live evals; if more than one cell is running, pick which feed "
        "to follow."
    )

    # Above the tabs because it scopes all of them: the dataset the instructions
    # point at, the policy the Live tab defaults to, and the board shown under
    # Leaderboard.
    task_sel = gr.Radio(
        config.TASK_LABELS,
        value=config.DEFAULT_TASK_LABEL,
        label="Task",
    )
    task_blurb = gr.Markdown(_task_blurb(_DEFAULT_TASK))

    with gr.Tabs():
        with gr.Tab("Instructions"):
            instructions = gr.Markdown(_instructions_markdown(_DEFAULT_TASK))
        with gr.Tab("Live"):
            if not config.eval_image_configured():
                gr.Markdown(
                    "> ⚠️ **Not configured.** Set `ARMNET_DEMO_EVAL_IMAGE` (a pre-built "
                    "lerobot-eval runtime image ref) in the Space settings — see the README."
                )

            # Not "cells offline": a page that never gets its first refresh
            # would then look exactly like a fleet that is down.
            cell_light = gr.HTML(_light_html("grey", "checking cells…"))

            with gr.Row():
                preset = gr.Dropdown(
                    label="Pre-trained policy",
                    choices=list(_DEFAULT_TASK.presets),
                    value=_DEFAULT_TASK.reference_label,
                )
                custom = gr.Textbox(
                    label="…or a custom HF repo id",
                    placeholder=_DEFAULT_TASK.reference_policy,
                )
                run_btn = gr.Button(
                    f"Run {config.N_EPISODES} rollouts",
                    variant="primary",
                    interactive=False,
                    scale=0,
                )

            banner = gr.Markdown("")
            queue_view = gr.Markdown(
                "### Shared queue (0 waiting)\n\n"
                f"_Armnet eval submissions: 0/{config.MAX_QUEUED_JOBS}._\n\n"
                "_No jobs waiting._"
            )

            feed_sel = gr.Radio(
                label="Live eval",
                choices=[],
                visible=False,
            )
            feed_caption = gr.Markdown("")

            with gr.Row():
                with gr.Column(scale=3):
                    video = gr.Image(
                        label="Live rollout",
                        height=320,
                        interactive=False,
                        streaming=True,
                    )
                with gr.Column(scale=2):
                    status = gr.Markdown("")
                    summary = gr.Markdown("")
                    rollouts = gr.Markdown("### Rollouts\n\n_No rollouts scored yet._")

            events = gr.Textbox(label="Live log", lines=8, interactive=False, autoscroll=True)

        with gr.Tab("Leaderboard") as leaderboard_tab:
            gr.Markdown(
                "Each task has **its own board** — the one shown follows the task "
                "selected above. Results are **pooled by policy repo + revision** "
                "across every run, and ranked by the **95% lower-bound** success rate "
                "by default — so running more rollouts only helps if the policy is "
                "genuinely good (spamming low-quality evals can't game it). Switch to "
                "*Average success rate* to rank by the raw mean."
            )
            with gr.Row():
                lb_sort = gr.Radio(
                    ["Lower bound (95%)", "Average success rate"],
                    value="Lower bound (95%)",
                    label="Sort by",
                    scale=3,
                )
                lb_refresh = gr.Button("Refresh", scale=0)
            lb_table = gr.Dataframe(
                headers=leaderboard.HEADERS,
                datatype=leaderboard.DATATYPES,
                interactive=False,
                wrap=True,
                label="Leaderboard",
                elem_id="leaderboard-table",
            )

    # Text is cheap enough to poll. Video uses one generator response for the
    # lifetime of the page, so increasing its FPS does not multiply HTTP/API
    # requests per viewer.
    text_timer = gr.Timer(_TEXT_REFRESH_SECONDS)

    text_outputs = [
        status,
        rollouts,
        summary,
        events,
        run_btn,
        cell_light,
        queue_view,
        feed_sel,
        feed_caption,
    ]

    run_btn.click(
        _on_run,
        inputs=[task_sel, preset, custom],
        outputs=[banner, status, run_btn, feed_sel, feed_caption, cell_light],
    )
    # concurrency_limit=None on both: they share one Gradio concurrency group
    # (keyed by the function), and its default limit of 1 would put every
    # viewer's refresh behind a single worker. It only reads a locked snapshot.
    text_timer.tick(
        _refresh_text,
        inputs=[task_sel, feed_sel],
        outputs=text_outputs,
        concurrency_limit=None,
    )
    feed_sel.change(
        _on_feed_change,
        inputs=[feed_sel, task_sel],
        outputs=[status, rollouts, summary, events, feed_caption],
    )

    demo.load(
        _refresh_text,
        inputs=[task_sel, feed_sel],
        outputs=text_outputs,
        concurrency_limit=None,
    )
    demo.load(
        _stream_video,
        outputs=video,
        show_progress="hidden",
        concurrency_limit=None,
        stream_every=_VIDEO_STREAM_SECONDS,
        api_visibility="private",
    )

    # Switching task re-points every tab at once, so a viewer never sees one
    # task's dataset next to another's leaderboard.
    task_sel.change(
        _on_task_change,
        inputs=[task_sel, lb_sort],
        outputs=[
            instructions,
            preset,
            custom,
            cell_light,
            run_btn,
            lb_table,
            task_blurb,
            feed_sel,
            feed_caption,
        ],
    )

    # Leaderboard: load on tab open / sort change / manual refresh (and once on
    # page load). Not on a timer — it changes only when a run finishes, so
    # constant polling would just risk rate limiting for no benefit.
    leaderboard_tab.select(_load_leaderboard, inputs=[task_sel, lb_sort], outputs=[lb_table])
    lb_sort.change(_load_leaderboard, inputs=[task_sel, lb_sort], outputs=[lb_table])
    lb_refresh.click(
        lambda t, s: _load_leaderboard(t, s, force=True),
        inputs=[task_sel, lb_sort],
        outputs=[lb_table],
    )
    demo.load(_load_leaderboard, inputs=[task_sel, lb_sort], outputs=[lb_table])


if __name__ == "__main__":
    demo.queue().launch()