"""armnet eval Space configuration.

A deliberately constrained public demo of the armnet infra: no install, no auth
setup, no Docker. A pre-provisioned API key is baked in (it is a throwaway demo
key, intended to be invalidated later), jobs route to any free BusyBox cell, and
the runtime image is pre-built and referenced by tag (a Space can't build Docker).

Everything that differs between the tasks on offer — the dataset to train
against, the reference policy, the leaderboard scope — hangs off :data:`TASKS`,
so adding a third task is a matter of one more entry rather than a sweep through
the UI.
"""

from __future__ import annotations

import os
from dataclasses import dataclass

# Demo API key. Provided as a Space secret (ARMNET_API_KEY) — nothing is baked
# into the repo, so no key ships in source. armnet-client reads it straight from
# the environment (armnet_client._config.api_key()).
DEMO_API_KEY = os.environ.get("ARMNET_API_KEY")

# Every task runs on the single-arm BusyBox cell. Routing is by task slug alone:
# the orchestrator resolves the slug to the `busybox` environment and sends the
# job to a cell serving it, so the Space never names a cell in a submission.
# Only tasks that auto-score off the instrumented box and reset themselves
# between rollouts belong here — one needing a human to tidy up between rollouts
# would stall an unattended public demo.
EMBODIMENT = "lerobot/so-101"
MULTITASK_DATASET = "armnet/busybox_multitask"


@dataclass(frozen=True)
class Task:
    """One evaluable task: what to train on, what to beat, what to rank."""

    slug: str
    """The armnet task slug submitted with the job."""

    label: str
    """What the selector calls it."""

    blurb: str
    """One line describing the rollout, shown under the selector."""

    dataset: str
    """HF dataset to train against, linked from the instructions."""

    reference_policy: str
    """Pre-trained policy offered as the default in the Live tab."""

    reference_label: str
    """How the reference policy is named in the dropdown."""

    @property
    def presets(self) -> dict[str, str]:
        return {self.reference_label: self.reference_policy}


TASKS: dict[str, Task] = {
    task.label: task
    for task in (
        Task(
            slug="push_green_button",
            label="Push the green button",
            blurb=(
                "Push the green button on the BusyBox panel. The button is "
                "momentary and springs back, so the scene resets itself."
            ),
            dataset="armnet/busybox_push_green_button",
            reference_policy="armnet/busybox_push_green_button_diffusion",
            reference_label="Diffusion · push green button — reference",
        ),
        Task(
            slug="busybox_left_toggle_down",
            label="Flip left switch off",
            blurb=(
                "Flip the left switch down into the off position. The arm flips it "
                "back up between rollouts, so the scene resets itself."
            ),
            dataset="armnet/busybox_flip_left_switch_off",
            reference_policy="armnet/busybox_flip_left_switch_off_diffusion",
            reference_label="Diffusion · flip left switch off — reference",
        ),
    )
}

TASK_LABELS = list(TASKS)
DEFAULT_TASK_LABEL = TASK_LABELS[0]


def task_for(label: str) -> Task:
    """The selected task, falling back to the default for an unknown label."""

    return TASKS.get(label) or TASKS[DEFAULT_TASK_LABEL]

# Fallback cell id written on leaderboard rows that predate per-job attribution
# (and on a run whose assigned cell is not yet known). Live availability no
# longer names a particular cell — the UI reports how many cells are free.
CELL_ID = os.environ.get("ARMNET_DEMO_CELL_ID", "cell-01")

# How often the background poller refreshes fleet availability. This is a
# single process-wide poll shared by every viewer (not per tab), so it can
# stay brisk without adding to the per-viewer request budget.
CELL_STATUS_POLL_S = float(os.environ.get("ARMNET_DEMO_CELL_STATUS_POLL_S", "5.0"))

# Number of rollouts per eval run. 20 to keep per-policy variance down.
N_EPISODES = int(os.environ.get("ARMNET_DEMO_N_EPISODES", "20"))
# Protect the shared public queue from unbounded submissions. The orchestrator
# Client-side abuse guard: only queued jobs submitted by this Space count.
MAX_QUEUED_JOBS = int(os.environ.get("ARMNET_DEMO_MAX_QUEUED_JOBS", "10"))
SUBMISSION_SOURCE = "armnet_eval_space"

# SDK-side wait budget. 10 rollouts x ~30s + resets + policy load; generous.
TIMEOUT_SECONDS = int(os.environ.get("ARMNET_DEMO_TIMEOUT_S", "2400"))

# Pre-built + pre-pushed runtime image ref. A Space can't build Docker, so the
# lerobot-eval runtime (which bakes in the molmoact2 plugin) is built and pushed
# elsewhere and named here. The policy itself is NOT baked in — it is loaded at
# runtime from the repo id passed as a job arg, so one image serves every policy.
#
# The ARMNET_DEMO_EVAL_IMAGE variable is what actually decides, and the
# armnet-demo-eval-image workflow keeps it pointing at the current runtime. This
# constant is only the fallback for a Space (or a local run) without it set, so
# treat it as a floor rather than the truth: it is bumped by hand and will lag.
_DEFAULT_EVAL_IMAGE = (
    "europe-north1-docker.pkg.dev/alpha-robotics-476817/customer-images/"
    "ville/armnet-lerobot-eval:armnet-01c6ad4641e6"
)
EVAL_IMAGE = os.environ.get("ARMNET_DEMO_EVAL_IMAGE") or _DEFAULT_EVAL_IMAGE

# Per-rollout wall-clock and control rate (mirror armnet-lerobot-eval defaults).
EPISODE_TIME_S = float(os.environ.get("ARMNET_DEMO_EPISODE_TIME_S", "30.0"))
POLICY_FPS = int(os.environ.get("ARMNET_DEMO_POLICY_FPS", "20"))
START_SEED = 1000

# Every leaderboard policy is scored across the same spread of rail positions,
# camera angles and lighting rather than one staged pose, so the ranking
# rewards generalisation. Not exposed to the submitter and deliberately not
# read from the environment: a seed anyone could change is a leaderboard whose
# entries are not comparable.
VARIATION_SEED = 42

# Optional HF token. Two uses: (1) the Space stores it once as a named
# orchestrator secret (EVAL_HF_SECRET_NAME) under the demo user and the eval
# references it as $HF_TOKEN, so the container can pull a gated policy; (2) it
# authenticates the leaderboard dataset writes. Set it as a Space secret; without
# it only fully-public policies load and the leaderboard is read-only.
HF_TOKEN = os.environ.get("HF_TOKEN")

# Name of the orchestrator secret (under the demo user) that holds the HF token.
# The job references it as {"HF_TOKEN": EVAL_HF_SECRET_NAME}; the orchestrator
# resolves the *name* to its value — it never accepts an inline token value.
EVAL_HF_SECRET_NAME = os.environ.get("ARMNET_DEMO_HF_SECRET_NAME", "hf_token")

# Push the recorded eval dataset to the Hub. Off for the demo: it avoids needing
# Hub write access and spamming datasets; the live results still show in the UI.
PUSH_DATASET = os.environ.get("ARMNET_DEMO_PUSH_DATASET", "").lower() in ("1", "true", "yes")

# Namespace for that dataset. The eval otherwise calls whoami() and uses the
# token owner's personal account, which is not the org the token can write to.
DATASET_NAMESPACE = os.environ.get("ARMNET_DEMO_DATASET_NAMESPACE", "armnet")

# Live-video stream size, passed to the eval as job args. Small values keep the
# NATS frames tiny so the stream stays live and doesn't drop frames (e.g. the
# success/fail moment) under backpressure. scale=0.3 -> ~9% of the pixels;
# jpeg_quality=30 -> strong compression. Together ~10% of the original size.
RERUN_IMAGE_SCALE = float(os.environ.get("ARMNET_DEMO_RERUN_IMAGE_SCALE", "0.3"))
RERUN_JPEG_QUALITY = int(os.environ.get("ARMNET_DEMO_RERUN_JPEG_QUALITY", "30"))


# Leaderboard persistence. Gradio has no built-in DB and Space disk is
# ephemeral, so completed runs are appended to this HF *dataset* repo (reloaded
# on restart). Reads are public; writes need HF_TOKEN (same one used for gated
# policy pulls). Pooled by (policy repo, revision).
LEADERBOARD_REPO = os.environ.get(
    "ARMNET_DEMO_LEADERBOARD_REPO", "armnet/armnet-demo-leaderboard"
)


def eval_image_configured() -> bool:
    return bool(EVAL_IMAGE)