controller.py

"""Run controller — cooperative flags, control handlers, and the streaming
generators behind Run / Push-to-branch / Retry.

The heavy work runs in a worker thread; its output (log() + every print()) streams
through a queue back to the generator, which re-renders the page on each tick. Pure
rendering lives in views.py; the episode-review panel lives in review.py.

Nothing about a run is module-global: every flag lives on the caller's
session.RunSession (resolved from the injected `gr.Request`), so two people using
the Space at the same time never see — or control — each other's run.
"""

import contextlib
import queue
import tempfile
import threading
import time
from pathlib import Path

import gradio as gr
from huggingface_hub import HfApi, snapshot_download
from huggingface_hub.utils import GatedRepoError, RepositoryNotFoundError

import session
from pipeline import build_lerobot, episode_label, push_lerobot
from views import bar, btns, error_card, inputs_view, io, run_recap, success_summary

# The cooperative control flags live on the caller's session (session.RunSession):
#   stop           — abandon the run (Cancel confirmed)
#   pause          — hold the worker at the next safe checkpoint (Pause)
#   cancel_pending — Cancel was clicked, awaiting confirmation
# The streaming generator is the ONLY writer of button state: it re-renders the
# buttons from its own session's flags on every tick. The handlers below just flip
# a flag (outputs=None), so they never fight the generator for the same components.
# `request: gr.Request` is injected by Gradio — it is NOT a UI input.


def toggle_pause(request: gr.Request = None):
    """Pause ⇄ Resume: flip this session's pause flag (button relabels itself on
    the generator's next tick)."""
    sess = session.get(request)
    sess.pause.clear() if sess.pause.is_set() else sess.pause.set()


def request_cancel(request: gr.Request = None):
    """Cancel clicked: reveal the confirm/keep buttons (handled by the generator)."""
    session.get(request).cancel_pending.set()


def keep_running(request: gr.Request = None):
    """'Keep running' clicked: dismiss the cancel confirmation."""
    session.get(request).cancel_pending.clear()


def confirm_cancel(request: gr.Request = None):
    """Confirm-cancel handler: only signal this session's worker to abandon
    (outputs=None, like the old Stop, so it preempts the running generator via
    cancels=). The UI reset is done by reset() chained as .then() on this click — a
    single function that both cancels= and writes outputs can't reliably apply its
    outputs (it'd queue behind the generator it's cancelling)."""
    session.get(request).stop.set()


class _LineQueueWriter:
    """A file-like object that pushes complete lines onto a queue.

    Used to redirect stdout: every print() from convert / SLAM / build_dataset
    becomes a streamed log line, not just the explicit log() callback.
    """

    def __init__(self, q: "queue.Queue"):
        self.q = q
        self._buf = ""

    def write(self, s: str) -> int:
        self._buf += s
        while "\n" in self._buf:
            line, self._buf = self._buf.split("\n", 1)
            if line.strip():
                self.q.put(("log", line))
        return len(s)

    def flush(self):
        if self._buf.strip():
            self.q.put(("log", self._buf))
        self._buf = ""


def _preflight(api, source_repo, target_repo, scopes=frozenset()):
    """Quick access/existence checks. Returns (exists, writable, error_or_None).

    error_or_None is a clear, user-facing message (never a raw HF traceback) when
    the source can't be read or the target can't be written — shown in the red
    error card before any heavy work starts. Both checks run up front so a bad
    source or an unwritable target fails fast, not after a long SLAM run.

    scopes: the OAuth scopes actually granted to the sign-in token. Owning the
    namespace isn't enough to *create* a repo there — HF gates creation behind the
    'manage-repos' scope (writing content to an existing repo only needs
    'write-repos'). So if the target doesn't exist yet and that scope wasn't
    granted (partial consent), creation would 403 *after* a long SLAM run; we
    catch it here instead.
    """
    # ---- Source: must exist and be readable with this token ----
    try:
        api.repo_info(source_repo, repo_type="dataset")
    except RepositoryNotFoundError:
        return None, None, (
            f"Error on the source dataset '{source_repo}':\nEither it doesn't exist, or it's private and "
            f"your account can't see it.\nCheck the spelling; it should look like "
            f"'username/dataset-name'."
        )
    except GatedRepoError:
        return None, None, (
            f"Error on the source dataset '{source_repo}':\nIt is gated. Accept its access terms on "
            f"the Hub first, then re-run."
        )
    except Exception as e:
        return None, None, f"Cannot access source dataset '{source_repo}': {e}"

    # ---- Target: resolve namespace + whether this account can write to it ----
    # whoami() returns each org with the user's role; only write-capable roles can
    # push datasets, so being *in* an org isn't enough — check the role too.
    WRITE_ROLES = {"admin", "write", "contributor"}
    try:
        me = api.whoami()
        username = me.get("name")
        org_roles = {o.get("name"): o.get("roleInGroup") for o in me.get("orgs", [])}
    except Exception:
        username, org_roles = None, {}
    ns = target_repo.split("/")[0] if "/" in target_repo else username

    if ns is None:
        writable = False
    elif ns == username:
        writable = True
    elif ns in org_roles:
        writable = org_roles[ns] in WRITE_ROLES
    else:
        writable = False

    exists = api.repo_exists(target_repo, repo_type="dataset")

    if not writable:
        if username is None:
            return exists, False, (
                "Couldn't confirm your Hugging Face identity from the sign-in "
                "token. Sign out and back in, then re-run."
            )
        verb = "push to" if exists else "create"
        if ns in org_roles:
            role = org_roles[ns] or "read-only"
            return exists, False, (
                f"Your role in the '{ns}' org is '{role}', which can't write "
                f"datasets : so you can't {verb} '{target_repo}'. \nAsk an org admin "
                f"for write access, or set the target to a namespace you own."
            )
        who = f"'{username}'" + (
            f" (orgs: {', '.join(sorted(org_roles))})" if org_roles else " (no orgs)")
        return exists, False, (
            f"You don't have write access to namespace '{ns}', so you can't "
            f"{verb} '{target_repo}'.\nYou're signed in as {who} : set the target "
            f"to a namespace you own."
        )

    # Owning/writing the namespace is necessary but not sufficient to *create* a
    # repo: creation needs the 'manage-repos' scope. Only enforce it when the
    # target doesn't exist yet — an existing repo is reached by writing content
    # ('write-repos'), and a non-branch existing target is intercepted upstream.
    if not exists and "manage-repos" not in scopes:
        return exists, writable, (
            f"'{target_repo}' doesn't exist yet, and creating a new dataset needs "
            f"the 'manage-repos' permission — which your sign-in is missing.\n"
            f"Sign out and back in to grant it, then re-run. (Pushing to an "
            f"existing dataset only needs write access, which you have.)"
        )

    return exists, writable, None


def _reset_view(sess):
    """Full reset to the idle page (clears log/summary/bar, only Run shown).

    Clears the UI-facing flags (pause, cancel-pending, review) but NOT stop: when
    this runs as the .then() after a confirmed Cancel, the background worker is still
    winding down and relies on stop staying set to abandon (skip the push). _run
    clears it itself at the start of the next run (sess.start_run()).
    """
    sess.pause.clear()
    sess.cancel_pending.clear()
    sess.reset_review()
    return (*io("idle", bar(0, "ready"), ""), "", None, *btns("idle"),
            *inputs_view("idle"))


def _run(source_repo, target_repo, task, run_sync, tactile, oauth_token, to_branch, sess):
    """Streaming generator behind Run / Push-to-branch. Yields the 16-element
    `outputs` tuple: bar, log, summary, retry_state, the 8 buttons, then the 4
    inputs-view fields. The episode-review panel is a separate reactive @gr.render
    block (review.build_panel), not part of this tuple.

    Pre-flight checks run synchronously up front (source read, target write, target
    existence — an existing target reveals the "push to a branch" button); the heavy
    work then runs in a worker thread and streams through a queue.
    """
    if oauth_token is None:
        raise gr.Error("Sign in with your Hugging Face account first.")
    # Trim stray whitespace from the textboxes — a leading/trailing space in a
    # repo id is an easy copy-paste mistake that would otherwise 404 the source.
    source_repo = (source_repo or "").strip()
    target_repo = (target_repo or "").strip()
    task = (task or "").strip()
    if not (source_repo and target_repo and task):
        raise gr.Error("Fill in source repo, target repo and task.")

    # The signed-in user's OAuth token stays a LOCAL variable, passed explicitly to
    # every HF call below (HfApi / snapshot_download / build_lerobot / push_lerobot).
    # It must never be written to os.environ: the process env is shared by every
    # session and by the REST jobs, so a global HF_TOKEN would leave this user's
    # credential as the process-wide default for everyone else's requests (and
    # huggingface_hub.login() wouldn't take an OAuth token anyway — it expects a
    # classic token's role). push_lerobot needs it in the env for LeRobot's
    # push_to_hub, and scopes it there under a lock for that call only.
    token = oauth_token.token
    scopes = frozenset((getattr(oauth_token, "scope", "") or "").split())

    sess.start_run()  # fresh run: clear this session's flags

    logs: list[str] = []
    start = time.time()
    frac, label = 0.02, "checking"
    retry_ctx = None  # set when a push fails but the built dataset is reusable

    def render() -> str:
        return "\n".join(logs)

    def view(summary="", *, state="running", branch=False, retry=False):
        in_review = sess.review.is_set()
        lbl = ("awaiting review" if in_review
               else f"paused · {label}" if sess.pause.is_set()
               else f"{label} · {int(time.time() - start)}s")
        return (*io(state, bar(frac, lbl), render()), summary, retry_ctx,
                *btns(state, paused=sess.pause.is_set(),
                      cancel_pending=sess.cancel_pending.is_set(),
                      branch=branch, retry=retry,
                      allow_pause=not in_review),
                *inputs_view(state, source=source_repo, target=target_repo,
                             task=task))

    # ---- Pre-flight (synchronous, fast) ---------------------------------
    logs.append("Checking repo access…")
    yield view()  # 'Running' + Pause + Cancel appear immediately
    api = HfApi(token=token)
    exists, _writable, err = _preflight(api, source_repo, target_repo, scopes)
    if err:
        logs.append(err)
        # Recoverable: back to idle so the user can fix the inputs and re-run.
        frac, label = 0.0, "error"
        yield view(summary=error_card(err), state="idle")
        return
    logs.append("  ✓ source readable")
    logs.append(f"  ✓ write access to '{target_repo.split('/')[0]}'")

    if not exists:
        logs.append(f"Target '{target_repo}' is new and writable")
    elif not to_branch:
        logs.append(f"Target '{target_repo}' already exists")
        warn = (
            f"### ⚠️ Target dataset already exists\n"
            f"`{target_repo}` already exists on the Hub.\nTo avoid overwriting it, "
            f"click **Push to a new branch** below: your result lands on a "
            f"`grabette-…` branch and `main` is left untouched.\n\n"
            f"_(Want a different target instead? Click **Reset** and change it.)_"
        )
        frac, label = 0.05, "target exists"
        # Recoverable: back to idle (re-run after editing the target) but reveal
        # the "push to a branch" button.
        yield view(summary=warn, state="idle", branch=True)
        return
    else:
        logs.append(f"  ✓ target '{target_repo}' exists — will push to a branch")
    yield view()

    # ---- Heavy work in a worker thread ----------------------------------
    q: "queue.Queue[tuple]" = queue.Queue()
    result: dict = {}

    def on_progress(done, total, phase):
        if phase == "check":
            f = 0.10 + 0.18 * (done / total if total else 1.0)
            lbl = f"checking {done}/{total}" if total else "checking"
        elif phase == "slam":
            f = 0.30 + 0.55 * (done / total if total else 1.0)
            lbl = f"SLAM {done}/{total}" if total else "SLAM"
        elif phase == "build":
            f, lbl = 0.88, "building dataset"
        elif phase == "push":
            f, lbl = 0.95, ("pushing branch" if to_branch else "pushing")
        else:
            f, lbl = 0.30, phase
        q.put(("progress", f, lbl))

    def _review_gate(kind, items, all_eps):
        """Publish `items` to THIS session's (kind="input"/"trajectory") review panel,
        block the worker until Continue (or a cancel), then return the kept episode
        dirs. Another user's panel never sees these episodes: they are published on
        the session captured by this run, not on a module-level flag.

        all_eps is the full [(name, ep_dir), …] in order, where name is the
        episode's UNIQUE label (episode_label — in the fleet's role layout the dir
        name alone is just "left"/"right", shared by every recording, so dropping
        one card would drop that arm of every recording); the kept subset is the
        episodes the user did NOT remove (🗑). On cancel the full set is returned
        and the build then hits should_stop and raises — so nothing is pushed."""
        sess.review_kind = kind
        sess.review_items[:] = items
        sess.review_drop.clear()
        sess.review_done.clear()
        sess.review.set()
        while not sess.review_done.is_set():
            if sess.stop.is_set():
                sess.review.clear()
                return [ep for _, ep in all_eps]
            time.sleep(0.2)
        sess.review.clear()
        drop = set(sess.review_drop)
        kept = [ep for name, ep in all_eps if name not in drop]
        if drop:
            print(f"Dropped {len(drop)} episode(s): {', '.join(sorted(drop))}.")
        print(f"Keeping {len(kept)} episode(s).")
        return kept

    def pre_review_cb(checks):
        """Called by build_lerobot between the completeness/sync prechecks and SLAM.
        checks is a list of {ep, name, errors, warnings, sync} dicts. Returns the
        episode dirs to keep — letting the user drop incomplete/desynced recordings
        BEFORE the slow SLAM runs. A clean run (nothing flagged) continues untouched."""
        def flagged(c):
            return bool(c["errors"] or c["warnings"]) or (
                c["sync"] is not None and c["sync"]["verdict"] != "GOOD")
        bad = [c for c in checks if flagged(c)]
        if not bad:
            return [c["ep"] for c in checks]
        items = [
            {"name": c["name"], "kind": "input",
             "verdict": ("ERROR" if c["errors"]
                         or (c["sync"] and c["sync"]["verdict"] == "BAD") else "WARN"),
             "messages": [
                 *(f"[error] {m}" for m in c["errors"]),
                 *c["warnings"],
                 *([c["sync"]["message"]]
                   if c["sync"] and c["sync"]["verdict"] != "GOOD" else []),
             ]}
            for c in bad
        ]
        print(f"⏸ Pre-SLAM review: {len(bad)} of {len(checks)} episode(s) flagged by "
              f"the dataset / sync check — remove the ones to drop (🗑), then Continue.")
        return _review_gate("input", items, [(c["name"], c["ep"]) for c in checks])

    def review_cb(results):
        """Called by build_lerobot between SLAM and the build. results is a list of
        (episode_dir, TrajectoryReport). Returns the episode dirs to keep — letting
        the user drop episodes whose trajectory came back flagged. A clean run
        (all GOOD) continues untouched."""
        flagged = [(ep, rep) for ep, rep in results if rep.verdict != "GOOD"]
        if not flagged:
            return [ep for ep, _ in results]
        items = [
            {"name": episode_label(ep), "verdict": rep.verdict,
             "n_tracked": rep.n_tracked, "tracking_pct": rep.tracking_pct,
             "total_distance_m": rep.total_distance_m, "duration_s": rep.duration_s,
             "median_step_mm": rep.median_step_mm, "median_angle_deg": rep.median_angle_deg,
             "n_jumps": rep.n_jumps,
             "messages": [*rep.errors, *rep.warnings]}
            for ep, rep in flagged
        ]
        print(f"⏸ Trajectory review: {len(flagged)} of {len(results)} episode(s) flagged "
              f"— remove the ones to drop (🗑), then click Continue.")
        return _review_gate("trajectory", items,
                            [(episode_label(ep), ep) for ep, _ in results])

    def worker():
        writer = _LineQueueWriter(q)
        try:
            with contextlib.redirect_stdout(writer):
                q.put(("progress", 0.07, "download"))
                work = Path(tempfile.mkdtemp())
                try:
                    info = api.repo_info(source_repo, repo_type="dataset", files_metadata=True)
                    mb = sum((s.size or 0) for s in info.siblings) / 1e6
                    print(f"Downloading {source_repo} — {len(info.siblings)} files, {mb:.0f} MB…")
                except Exception:
                    print(f"Downloading {source_repo}…")
                # Raw episodes are dominated by many small depth PNGs (≈600/ep),
                # so the download is bound by per-file request overhead, not
                # bandwidth — more concurrent connections is the main lever.
                raw = snapshot_download(source_repo, repo_type="dataset",
                                        local_dir=work / "raw", token=token,
                                        max_workers=48)
                print("Download complete.\n\n")
                q.put(("progress", 0.30, "converting"))
                ds_root = work / "lerobot"
                # One entry per LeRobot episode written — in a bimanual raw
                # dataset that's one per recording (both arms), not one per arm.
                built = build_lerobot(
                    raw, target_repo, task, root=ds_root,
                    log=print, should_stop=sess.stop.is_set, to_branch=to_branch,
                    on_progress=on_progress, token=token, gate=sess.wait_if_paused,
                    pre_review=pre_review_cb, review=review_cb, run_sync=run_sync,
                    tactile=tactile)
                # Build done — dataset cached on disk; record it so that if the push
                # fails, the "Retry push" button can reuse it (no re-running SLAM).
                result["built"] = {"root": str(ds_root), "n": len(built)}
                n, link, mode = push_lerobot(
                    target_repo, task, ds_root, len(built),
                    to_branch=to_branch, token=token, log=print,
                    on_progress=on_progress, gate=sess.wait_if_paused)
                writer.flush()
            result["n"], result["link"], result["mode"] = n, link, mode
        except Exception as e:
            writer.flush()
            result["error"] = e
        finally:
            q.put(("done",))

    t = threading.Thread(target=worker, daemon=True)
    t.start()

    while True:
        try:
            # Short timeout so Pause/Cancel button changes reflect within a tick.
            item = q.get(timeout=0.3)
        except queue.Empty:
            yield view()  # tick the elapsed counter / reflect pause+cancel flags
            continue
        tag = item[0]
        if tag == "done":
            break
        if tag == "progress":
            frac, label = item[1], item[2]
        else:  # ("log", line)
            logs.append(item[1])
        yield view()

    t.join()
    if "error" in result:
        err = result["error"]
        if result.get("built"):
            # Build succeeded, push didn't — offer a push-only retry, no re-SLAM.
            retry_ctx = {"target_repo": target_repo, "task": task, "to_branch": to_branch,
                         "root": result["built"]["root"], "n": result["built"]["n"]}
            msg = error_card(
                f"Push failed: {err}\n\nThe dataset is built and cached — click "
                f"“Retry push” to push it again without re-running SLAM.")
            label = "push failed"
            yield view(summary=msg + run_recap(logs), state="finished", retry=True)
        else:
            label = "failed"
            yield view(summary=error_card(f"Pipeline failed: {err}") + run_recap(logs),
                       state="finished")
        return

    n = result["n"]
    link = result.get("link")
    mode = result.get("mode")
    logs.append(f"✅ Done — {n} episode(s).")
    frac, label = 1.0, "done"
    yield view(summary=success_summary(target_repo, n, link, mode) + run_recap(logs),
               state="finished")


def retry_push(retry_ctx, oauth_token: gr.OAuthToken | None = None,
               request: gr.Request = None):
    """Retry just the push of an already-built dataset — no re-download, no SLAM.
    Reuses the on-disk dataset captured in retry_ctx by a previous failed run."""
    if oauth_token is None:
        raise gr.Error("Sign in with your Hugging Face account first.")
    if not retry_ctx:
        raise gr.Error("Nothing to retry — run the pipeline first.")
    token = oauth_token.token  # local only — never os.environ, see _run

    sess = session.get(request)
    sess.start_run()

    logs: list[str] = []
    start = time.time()
    frac, label = 0.95, "pushing"

    def render() -> str:
        return "\n".join(logs)

    def view(summary="", *, state="running", retry=False, ctx=retry_ctx):
        # 16-output shape, same as _run. A push can't be paused, so allow_pause
        # is False; Cancel still works (abandons the retry).
        return (*io(state, bar(frac, f"{label} · {int(time.time() - start)}s"), render()),
                summary, ctx,
                *btns(state, cancel_pending=sess.cancel_pending.is_set(),
                      retry=retry, allow_pause=False),
                *inputs_view(state, target=retry_ctx["target_repo"],
                             task=retry_ctx["task"]))

    logs.append("Retrying push (dataset already built — skipping SLAM)…")
    yield view()  # 'Running' + Cancel

    q: "queue.Queue[tuple]" = queue.Queue()
    result: dict = {}

    def on_progress(done, total, phase):
        q.put(("progress", 0.95, "pushing branch" if retry_ctx["to_branch"] else "pushing"))

    def worker():
        writer = _LineQueueWriter(q)
        try:
            with contextlib.redirect_stdout(writer):
                n, link, mode = push_lerobot(
                    retry_ctx["target_repo"], retry_ctx["task"], Path(retry_ctx["root"]),
                    retry_ctx["n"], to_branch=retry_ctx["to_branch"], token=token,
                    log=print, on_progress=on_progress)
                writer.flush()
            result["n"], result["link"], result["mode"] = n, link, mode
        except Exception as e:
            writer.flush()
            result["error"] = e
        finally:
            q.put(("done",))

    t = threading.Thread(target=worker, daemon=True)
    t.start()
    while True:
        try:
            item = q.get(timeout=0.3)
        except queue.Empty:
            yield view()
            continue
        if item[0] == "done":
            break
        if item[0] == "progress":
            frac, label = item[1], item[2]
        else:
            logs.append(item[1])
        yield view()
    t.join()

    if "error" in result:
        err = result["error"]
        # Keep the dataset cached: finished state with Reset, plus Retry to try
        # the push again.
        label = "push failed"
        yield view(summary=error_card(f"Push failed again: {err}"),
                   state="finished", retry=True, ctx=retry_ctx)
        return

    n = result["n"]
    link = result.get("link")
    mode = result.get("mode")
    logs.append(f"✅ Pushed — {n} episode(s).")
    frac, label = 1.0, "done"
    yield view(summary=success_summary(retry_ctx["target_repo"], n, link, mode),
               state="finished", ctx=None)


def run_pipeline(source_repo, target_repo, task, run_sync=False, tactile=True,
                 oauth_token: gr.OAuthToken | None = None,
                 request: gr.Request = None):
    yield from _run(source_repo, target_repo, task, run_sync, tactile, oauth_token,
                    to_branch=False, sess=session.get(request))


def run_pipeline_branch(source_repo, target_repo, task, run_sync=False, tactile=True,
                        oauth_token: gr.OAuthToken | None = None,
                        request: gr.Request = None):
    yield from _run(source_repo, target_repo, task, run_sync, tactile, oauth_token,
                    to_branch=True, sess=session.get(request))


def reset(request: gr.Request = None):
    """Reset button (shown when finished): clear the page and return to idle
    (only Run). Inputs (source/target/task) are kept so the user can re-run."""
    return _reset_view(session.get(request))