"""Metadata-only estimates and authenticated HF Jobs submission. No ambient credentials."""

import hashlib
import json
import math
import os
import re
import shlex
import tempfile
import time
import uuid
from pathlib import Path, PurePosixPath

import pyarrow.parquet as pq
from flora_captioner import FLORA
from flora_image import validated_worker_image
from hand_actions import COMPUTE, HANDS, pin_runtime, prepare_hands, runtime_configuration
from huggingface_hub import HfApi, hf_hub_download
from huggingface_hub.utils import validate_repo_id
from model_catalog import (
    DEFAULT_MODEL,
    HUMAN_VIDEO_PARALLEL_CHOICES,
    HUMAN_VIDEO_SEGMENT_CHOICES,
    human_video_profile,
    human_video_segment_cost,
    human_video_segment_reservation,
    model_profile,
)
from prompt_config import TEMPLATES, validate_prompts

MODEL = model_profile()["route"]
PROVIDER = model_profile()["provider"]
FEATURES = {
    "Subtasks": "subtask",
    "Plans": "plan",
    "Memory": "memory",
    "Task rephrasings": "task_aug",
    "Interjections & speech": "interjection",
    "Visual questions": "vqa",
    # Not a language style: videos are saved under annotation_studio/human_videos/.
    "Human demo videos": "human_video",
    HANDS: "hands",
}
HUMAN_VIDEOS = "Human demo videos"
COPY = "Annotated dataset copy"
EXPORT = "Annotation files only"  # Historical reports only; new runs always write Parquet.


def needs_provider(features, prompts=None, captioner="selected-vlm"):
    """Native hand-only and Flora workflows never need a hosted helper."""
    features = set(features) - {HANDS}
    if not features:
        return False
    return (
        captioner != "flora"
        or bool(set(features) - {"Subtasks", "Plans"})
        or bool((prompts or {}).get("plan_render"))
    )


def repo_id(value):
    value = value.strip()
    validate_repo_id(value)
    if value.count("/") != 1:
        raise ValueError("Use a dataset ID in the form owner/dataset.")
    return value


def relative_path(value):
    path = PurePosixPath(value)
    if path.is_absolute() or ".." in path.parts or "\\" in value:
        raise ValueError("Dataset metadata contains an unsafe file path.")
    return str(path)


def whole_number(value, name, minimum=0):
    number = float(value)
    if not math.isfinite(number) or not number.is_integer() or number < minimum:
        raise ValueError(f"{name} must be a whole number of at least {minimum}.")
    return int(number)


def video_cameras(info):
    """The RGB video streams the worker can actually decode."""
    if not str(info.get("codebase_version", "")).startswith("v3"):
        raise ValueError("This Space currently supports LeRobot v3 video datasets.")
    fps = float(info["fps"])
    if not math.isfinite(fps) or fps <= 0:
        raise ValueError("The dataset has no valid frame rate.")
    cameras = [
        k for k, v in info["features"].items() if v.get("dtype") == "video" and "depth" not in k.lower()
    ]
    if not cameras:
        raise ValueError("The dataset needs at least one supported RGB video camera.")
    return cameras


def read_cameras(dataset, token=False):
    """Discover views from info.json only, without episode tables or video downloads."""
    dataset = repo_id(dataset)
    hub_info = HfApi(token=token or False).dataset_info(dataset)
    with tempfile.TemporaryDirectory(prefix="annotation-cameras-") as cache:
        path = hf_hub_download(
            dataset,
            "meta/info.json",
            repo_type="dataset",
            revision=hub_info.sha,
            token=token or False,
            cache_dir=cache,
        )
        info = json.loads(Path(path).read_text())
    return {"dataset": dataset, "revision": hub_info.sha, "cameras": video_cameras(info)}


def read_selection(dataset, start, count, token=False, revision=None, cache_dir=None):
    """Read episode metadata, never videos or frame data, for a pinned v3 source."""
    dataset = repo_id(dataset)
    start = whole_number(start, "Starting episode")
    count = whole_number(count, "Episode count")
    api = HfApi(token=token or False)
    hub_info = api.dataset_info(dataset, revision=revision)
    revision = hub_info.sha
    with tempfile.TemporaryDirectory(prefix="annotation-metadata-") as temp:
        cache = cache_dir or temp

        def download(name):
            return hf_hub_download(
                dataset,
                relative_path(name),
                repo_type="dataset",
                revision=revision,
                token=token or False,
                cache_dir=cache,
            )

        info = json.loads(Path(download("meta/info.json")).read_text())
        cameras = video_cameras(info)
        fps = float(info["fps"])
        files = sorted(
            x.path
            for x in api.list_repo_tree(
                dataset, path_in_repo="meta/episodes", recursive=True, repo_type="dataset", revision=revision
            )
            if x.path.endswith(".parquet")
        )
        selected = []
        for name in files:
            table = pq.read_table(download(name))
            keep = [
                c
                for c in table.column_names
                if c in ("episode_index", "length") or c.startswith("data/") or c.startswith("videos/")
            ]
            for row in table.select(keep).to_pylist():
                if int(row["episode_index"]) >= start:
                    if int(row["length"]) <= 0:
                        raise ValueError("Episode metadata contains an empty episode.")
                    selected.append(row)
            # v3 episode metadata is stored in ascending episode order.
            if count and len(selected) >= count:
                break
        selected.sort(key=lambda r: r["episode_index"])
        if count:
            selected = selected[:count]
        if not selected or len({r["episode_index"] for r in selected}) != len(selected):
            raise ValueError("No episodes found, or duplicate episode IDs in metadata.")
        return {
            "dataset": dataset,
            "revision": revision,
            "info": info,
            "episodes": selected,
            "cameras": cameras,
            "seconds": sum(int(r["length"]) / fps for r in selected),
            "total_episodes": int(info["total_episodes"]),
            "start": start,
            "count": count,
            "private": bool(hub_info.private),
        }


def estimate(
    selection,
    features,
    vqa_hz,
    hours,
    cpu_rate,
    prompts=None,
    model_key=DEFAULT_MODEL,
    human_video_segments=None,
    human_video_parallel=None,
    captioner="selected-vlm",
):
    if captioner not in {"selected-vlm", "flora"}:
        raise ValueError("Choose a supported subtask captioner.")
    profile = model_profile(model_key)
    prompts = validate_prompts(prompts or {}, features)
    if not features or any(f not in FEATURES for f in features):
        raise ValueError("Choose at least one annotation type.")
    if HUMAN_VIDEOS in features:
        if "Subtasks" not in features:
            raise ValueError("Human demo videos are generated per subtask; also select Subtasks.")
        if human_video_segments not in HUMAN_VIDEO_SEGMENT_CHOICES:
            raise ValueError("Choose a supported number of videos per episode.")
        if human_video_parallel not in HUMAN_VIDEO_PARALLEL_CHOICES:
            raise ValueError("Choose a supported number of parallel videos.")
    if hours not in (1, 4, 12) or vqa_hz not in (0.1, 0.5, 1.0):
        raise ValueError("Choose a supported timeout and visual-question frequency.")
    n, seconds = len(selection["episodes"]), selection["seconds"]
    needs_subtasks = bool(set(features) - {"Visual questions", HANDS})
    if captioner == "flora" and not needs_subtasks:
        raise ValueError("Flora needs an annotation type that uses subtasks.")
    if captioner == "flora" and "plan_subtasks" in prompts:
        raise ValueError("Flora uses NVIDIA's native prompts. Reset the custom Subtasks prompt to use Flora.")
    if captioner == "flora" and selection["info"]["fps"] < 4:
        raise ValueError("Flora temporal segmentation needs at least 4 source FPS.")
    windows = sum(
        max(1, math.ceil(r["length"] / selection["info"]["fps"] / 600)) for r in selection["episodes"]
    )
    # Token estimates, not measured usage. Plans are deterministic and add no model calls.
    provider_subtasks = needs_subtasks and captioner != "flora"
    inputs = (seconds / 10 * 1500 + windows * 1500) if provider_subtasks else 0
    outputs = (seconds / 5 * 55 + windows * 120) if provider_subtasks else 0
    calls = windows if provider_subtasks else 0
    if "Memory" in features:
        memories = max(0, seconds / 5 - n)
        inputs += memories * 1800
        outputs += memories * 150
        calls += memories
    if "Task rephrasings" in features:
        inputs += n * 400
        outputs += n * 500
        calls += n
    if "Interjections & speech" in features:
        inputs += n * 3 * 4500
        outputs += n * 3 * 250
        calls += n * 3
    if "Visual questions" in features:
        ticks = seconds * vqa_hz + n
        inputs += ticks * 1500
        outputs += ticks * 200
        calls += ticks
    videos = videos_max = media = media_max = video_minutes = 0
    if HUMAN_VIDEOS in features:
        video_profile = human_video_profile()
        # One video per subtask. Subtasks grow with each episode's length (calibrated rate in
        # model_catalog.HUMAN_VIDEO); the upper bound allows 1.5x that; the per-episode cap applies.
        cap = human_video_segments or math.inf
        base, per = video_profile["subtasks_per_episode"], video_profile["seconds_per_subtask"]
        fps = selection["info"]["fps"]
        expected = [base + r["length"] / fps / per for r in selection["episodes"]]
        videos = sum(min(cap, e) for e in expected)
        videos_max = sum(min(cap, math.ceil(1.5 * e)) for e in expected)
        media = videos * human_video_segment_cost()
        media_max = videos_max * human_video_segment_cost()
        video_minutes = math.ceil(videos / human_video_parallel) * video_profile["minutes_per_video"]
        # Two VLM calls per video: the plan and edit prompt (6 thumbnails), then the video prompt
        # (edited frame + 6 thumbnails). The plan call runs for every capped subtask.
        inputs += videos * video_profile["vlm_input_tokens"]
        outputs += videos * video_profile["vlm_output_tokens"]
        calls += videos * 2
    if "plan_render" in prompts:
        plan_calls = max(n, seconds / 5) + (n * 3 if "Interjections & speech" in features else 0)
        inputs += plan_calls * 1500
        outputs += plan_calls * 250
        calls += plan_calls
    # Longer custom templates add input cost. Generation length remains uncertain.
    prompt_calls = {
        "plan_subtasks": windows,
        "plan_memory": max(0, seconds / 5 - n),
        "plan_task_rephrasings": n,
        "interjections_initial_speech": n,
        "interjections_interjection": n * 3,
        "vqa": seconds * vqa_hz + n,
        "plan_render": max(n, seconds / 5) + n * 3,
    }
    for key, text in prompts.items():
        extra = max(0, len(text.encode()) - len(TEMPLATES[key]["default"].encode()))
        inputs += extra * prompt_calls[key]
    inference = (inputs * profile["input_rate"] + outputs * profile["output_rate"]) / 1e6
    cpu_max = math.ceil(hours * 60 + 1) * cpu_rate
    # Range includes retries, grounding, variable segment counts and recovery.
    return {
        "low": inference * 0.6 + media * 0.8 + cpu_max,
        "high": inference * 2.5 + media_max + cpu_max,
        "inference": inference,
        "media": media,
        "media_max": media_max,
        "videos": round(videos),
        "videos_max": videos_max,
        "video_minutes": video_minutes,
        "video_parallel": human_video_parallel if videos else 0,
        "segment_usd": human_video_segment_cost() if videos else 0,
        "cpu_max": cpu_max,
        "cpu_rate": cpu_rate,
        "compute_flavor": FLORA["hardware"] if captioner == "flora" else "cpu-upgrade",
        "compute_name": "H200 GPU" if captioner == "flora" else "CPU",
        "helper_required": captioner == "flora" and needs_provider(features, prompts, captioner),
        "calls": math.ceil(calls),
        "custom_prompts": len(prompts),
        "hours": hours,
        "input_rate": profile["input_rate"],
        "output_rate": profile["output_rate"],
        "cached_input_rate": profile["cached_input_rate"],
        "price_date": profile["checked_at"],
        "model": profile["route"],
        "model_key": profile["key"],
        "context_length": profile["context_length"],
        "pricing_tiers": profile.get("pricing_tiers", []),
    }


def fingerprint(
    dataset,
    start,
    count,
    features,
    camera,
    vqa_hz,
    hours,
    mode,
    model_key=DEFAULT_MODEL,
    human_video_segments=None,
    human_video_parallel=None,
    captioner="selected-vlm",
):
    return hashlib.sha256(
        json.dumps(
            [
                dataset.strip(),
                start,
                count,
                sorted(features),
                camera,
                vqa_hz,
                hours,
                mode,
                model_key,
                human_video_segments,
                human_video_parallel,
                captioner,
            ],
            sort_keys=True,
        ).encode()
    ).hexdigest()


def validate_subset_features(info):
    for name, feature in info["features"].items():
        if feature["dtype"] == "video" and feature.get("info", {}).get("video.is_depth_map"):
            raise ValueError(f"Annotated datasets do not yet support depth video ({name}).")


def prepare(
    dataset,
    start,
    count,
    features,
    camera,
    vqa_hz,
    hours,
    mode,
    token=False,
    prompts=None,
    model_key=DEFAULT_MODEL,
    human_video_segments=None,
    human_video_parallel=None,
    captioner="selected-vlm",
    hand_actions=None,
):
    profile = model_profile(model_key)
    prompts = validate_prompts(prompts or {}, features)
    if HUMAN_VIDEOS not in features:
        human_video_segments = human_video_parallel = None
    if mode != COPY:
        raise ValueError("Annotations must be saved in a LeRobot Parquet dataset. Refresh the Space.")
    selection = read_selection(dataset, start, count, token)
    validate_subset_features(selection["info"])
    camera = camera or selection["cameras"][0]
    if camera not in selection["cameras"]:
        raise ValueError("Choose a camera in this dataset.")
    if captioner not in {"selected-vlm", "flora"}:
        raise ValueError("Choose a supported subtask captioner.")
    hands = prepare_hands(selection, camera, hand_actions, token) if HANDS in features else None
    if hands and captioner == "flora":
        raise ValueError("Run Hand actions and Flora separately so each model has the GPU memory it requires")
    flavor = hands["compute"] if hands else FLORA["hardware"] if captioner == "flora" else "cpu-upgrade"
    hardware = next(h for h in HfApi(token=token or False).list_jobs_hardware() if h.name == flavor)
    if hardware.unit_label != "minute":
        raise ValueError("HF changed its billing units; refresh the Space before launching.")
    costs = estimate(
        selection,
        features,
        vqa_hz,
        hours,
        hardware.unit_cost_usd,
        prompts,
        model_key,
        human_video_segments,
        human_video_parallel,
        captioner,
    )
    if hands:
        costs.update(compute_flavor=flavor, compute_name=COMPUTE[flavor], hand_timing=hands["timing"])
    return {
        "selection": selection,
        "features": list(features),
        "camera": camera,
        "vqa_hz": vqa_hz,
        "mode": mode,
        "prompts": prompts,
        "costs": costs,
        "model_profile": profile,
        "captioner": captioner,
        "captioner_profile": dict(FLORA) if captioner == "flora" else None,
        "hand_actions": hands,
        "human_video": {
            # 0 = every subtask.
            "segments_per_episode": human_video_segments,
            "parallel": human_video_parallel,
            "profile": human_video_profile(),
        }
        if HUMAN_VIDEOS in features
        else None,
        "created": time.time(),
        "nonce": uuid.uuid4().hex,
        "fingerprint": fingerprint(
            dataset,
            start,
            count,
            features,
            camera,
            vqa_hz,
            hours,
            mode,
            model_key,
            human_video_segments,
            human_video_parallel,
            captioner,
        ),
    }


def provider_live(api, model_id, provider):
    mapping = api.model_info(model_id, expand=["inferenceProviderMapping"]).inference_provider_mapping
    hosted = (
        mapping.get(provider)
        if isinstance(mapping, dict)
        else next((item for item in mapping or [] if item.provider == provider), None)
    )
    return hosted is not None and hosted.status == "live"


def job_command(space, revision, timeout_minutes=None, captioner="selected-vlm", hands=False):
    # Only trusted deployment values enter shell syntax. User options travel in JSON env.
    validate_repo_id(space)
    if not re.fullmatch(r"[a-f0-9]{40}", revision):
        raise ValueError("A pinned Space revision is required.")
    download = (
        "from huggingface_hub import snapshot_download; "
        f"snapshot_download({space!r},repo_type='space',revision={revision!r},"
        "local_dir='/workspace/app',allow_patterns=['*.py','*.txt','prompt-templates.json','hand-worker-image.json','pipeline-*'])"
    )
    return [
        "bash",
        "-lc",
        "\n".join(
            [
                "set -euo pipefail",
                "export HF_HUB_DISABLE_XET=1 HF_HUB_DISABLE_PROGRESS_BARS=1",
                # Keep native Hub transfers and byte progress for large hand-model assets.
                "export HF_HUB_DISABLE_XET=0 HF_HUB_DISABLE_PROGRESS_BARS=0 HF_HUB_DOWNLOAD_TIMEOUT=60 HF_HUB_ETAG_TIMEOUT=60"
                if hands
                else ":",
                f"export ANNOTATION_DEADLINE_EPOCH=$(( $(date +%s) + {timeout_minutes * 60 - 30} ))"
                if timeout_minutes
                else ":",
                "mkdir -p /workspace/app /workspace/source",
                "cd /workspace",
                "export PATH=/opt/hand-controller/bin:$PATH" if hands else ":",
                "export NVIDIA_TF32_OVERRIDE=0" if hands else ":",
                ":" if hands else "pip install --quiet huggingface_hub==1.26.0",
                "python -c " + shlex.quote(download),
                ":"
                if hands
                else "pip install --quiet torch==2.11.0 torchvision==0.26.0 --index-url https://download.pytorch.org/whl/"
                + ("cu128" if captioner == "flora" else "cpu"),
                ":" if hands else "pip install --quiet -r /workspace/app/worker-requirements.txt",
                "pip install --quiet -r /workspace/app/flora-requirements.txt"
                if captioner == "flora"
                else ":",
                "python /workspace/app/bootstrap.py",
                "export PYTHONPATH=/workspace/source/src:/workspace/app",
                "python /workspace/app/worker.py",
            ]
        ),
    ]


def launch(
    plan,
    output,
    budget,
    acknowledged,
    token,
    space=None,
    api_factory=HfApi,
    timeout_minutes=None,
    visibility="private",
):
    if visibility not in ("private", "public"):
        raise ValueError("Choose private or public output visibility.")
    hands = plan.get("hand_actions") if plan else None
    if hands and visibility != "private":
        raise ValueError("Hand research outputs must be private")
    if hands and hands["runtime"] != pin_runtime(runtime_configuration()):
        raise ValueError("The hand runtime changed; calculate a fresh estimate")
    if timeout_minutes is not None and (
        type(timeout_minutes) is not int or not 1 <= timeout_minutes <= plan["costs"]["hours"] * 60
    ):
        raise ValueError("Runtime limit must be within the quoted timeout.")
    if not token:
        raise ValueError("Sign in with Hugging Face before starting a job.")
    if not acknowledged:
        raise ValueError("Confirm the billing and data-sharing notice first.")
    if not plan or time.time() - plan["created"] > 1800:
        raise ValueError("Calculate a fresh estimate before starting a job.")
    if plan["mode"] != COPY:
        raise ValueError(
            "Annotations must be saved in a LeRobot Parquet dataset. Calculate a fresh estimate."
        )
    budget = float(budget)
    if not math.isfinite(budget) or not 1 <= budget <= 5000:
        raise ValueError("Choose a total spending target between $1 and $5,000.")
    if budget <= plan["costs"]["cpu_max"] + 0.05:
        raise ValueError("The spending target must cover the compute reservation and inference.")
    api = api_factory(token=token)
    username = api.whoami()["name"]
    output = repo_id(output)
    if output.split("/")[0] != username:
        raise ValueError("Use a new dataset in your personal Hugging Face account.")
    if output == plan["selection"]["dataset"]:
        raise ValueError("The output must be different from the source dataset.")
    # A repeated click / uncertain network reply must not create a second paid job.
    for job in api.list_jobs(namespace=username):
        if (getattr(job, "labels", None) or {}).get("annotation_request") == plan["nonce"]:
            return {"id": job.id, "owner": username, "output": output}
    if api.repo_exists(output, repo_type="dataset"):
        raise ValueError("That output dataset already exists. Choose a new name.")
    current = api.dataset_info(plan["selection"]["dataset"])
    if current.sha != plan["selection"]["revision"]:
        raise ValueError("The source changed. Calculate a fresh estimate before launching.")
    captioner = plan.get("captioner", "selected-vlm")
    if captioner not in {"selected-vlm", "flora"}:
        raise ValueError("Choose a supported subtask captioner.")
    if captioner == "flora" and plan.get("captioner_profile") != FLORA:
        raise ValueError("Flora configuration changed. Calculate a fresh estimate.")
    flavor = hands["compute"] if hands else FLORA["hardware"] if captioner == "flora" else "cpu-upgrade"
    hardware = next(h for h in api.list_jobs_hardware() if h.name == flavor)
    if hardware.unit_label != "minute" or hardware.unit_cost_usd > plan["costs"]["cpu_rate"]:
        raise ValueError("Compute pricing changed. Calculate a fresh estimate.")
    profile = model_profile(plan["costs"].get("model_key", DEFAULT_MODEL))
    if any(
        plan["costs"].get(k) != profile[k]
        for k in ("input_rate", "output_rate", "cached_input_rate", "context_length")
    ):
        raise ValueError("Model pricing changed. Calculate a fresh estimate.")
    if plan["costs"].get("model") != profile["route"]:
        raise ValueError("Model selection changed. Calculate a fresh estimate.")
    if plan["costs"].get("pricing_tiers", []) != profile.get("pricing_tiers", []):
        raise ValueError("Model pricing changed. Calculate a fresh estimate.")
    uses_provider = needs_provider(plan["features"], plan.get("prompts"), captioner)
    minimum_reservation = (
        profile["context_length"]
        * max([profile["input_rate"]] + [t["input_rate"] for t in profile.get("pricing_tiers", [])])
        + 4096 * max([profile["output_rate"]] + [t["output_rate"] for t in profile.get("pricing_tiers", [])])
    ) / 1e6
    if not uses_provider:
        minimum_reservation = 0
    human_video = plan.get("human_video")
    if human_video:
        if human_video["profile"] != human_video_profile():
            raise ValueError("Video pricing changed. Calculate a fresh estimate.")
        # One video must fit next to one conservative text reservation.
        minimum_reservation += human_video_segment_reservation()
    if budget - plan["costs"]["cpu_max"] < minimum_reservation:
        raise ValueError(
            f"This model needs at least ${minimum_reservation + plan['costs']['cpu_max']:.2f} for one conservative request reservation. Increase the target or choose a lower-cost model."
        )
    if uses_provider and not provider_live(api, profile["model_id"], profile["provider"]):
        raise ValueError("The selected model is temporarily unavailable from this provider.")
    if human_video and not all(
        provider_live(api, human_video["profile"][key], human_video["profile"]["provider"])
        for key in ("edit_model", "video_model")
    ):
        raise ValueError("The video models are temporarily unavailable from fal.")
    space = space or os.environ.get("SPACE_ID")
    if not space:
        raise ValueError("Job launch is enabled only on the deployed Hugging Face Space.")
    revision = api.space_info(space).sha
    manifest = {k: v for k, v in plan.items() if k not in ("selection", "fingerprint")}
    manifest.update(
        {
            "dataset": current.id,
            "dataset_revision": current.sha,
            "output": output,
            "visibility": visibility,
            "copy_scope": "annotated_episodes",
            "source_license": getattr(current, "card_data", {}).get("license")
            if getattr(current, "card_data", None)
            else None,
            "episode_ids_sha256": plan["selection"].get("episode_ids_sha256")
            or hashlib.sha256(
                json.dumps([r["episode_index"] for r in plan["selection"]["episodes"]]).encode()
            ).hexdigest(),
            "start": plan["selection"]["start"],
            "count": plan["selection"]["count"],
            "budget": budget,
            "inference_budget": budget - plan["costs"]["cpu_max"],
            "model": "EgoHand/SOMA + DROID-W + MoGe"
            if set(plan["features"]) == {HANDS}
            else profile["route"],
            "provider_required": uses_provider,
            "subtask_models": []
            if set(plan["features"]) == {HANDS}
            else [FLORA["segmenter_id"], FLORA["model_id"]]
            if captioner == "flora"
            else [profile["model_id"]],
            "helper_model": profile["route"] if captioner == "flora" and uses_provider else None,
            "model_profile": profile,
            "generation_options": profile["generation_options"],
            "request_concurrency": profile["concurrency"],
            "episode_parallelism": 2 if hands else 1 if captioner == "flora" else profile["concurrency"],
            "space": space,
            "space_revision": revision,
            "owner": username,
            "runtime_limit_minutes": timeout_minutes,
        }
    )
    image, command = "python:3.12-slim", job_command(space, revision, timeout_minutes, captioner)
    environment = {"ANNOTATION_MANIFEST": json.dumps(manifest)}
    accelerated = validated_worker_image(captioner, plan["features"])
    if accelerated:
        image_info = api.space_info(accelerated["repo"])
        integrated = accelerated.get("image_source") == "studio"
        if integrated:
            if accelerated["repo"] != space:
                raise ValueError("Integrated Flora must use this Annotation Studio Space")
            if image_info.runtime.stage != "RUNNING" or image_info.runtime.raw.get("sha") != image_info.sha:
                raise ValueError("Studio's CUDA image is still building; wait before starting Flora")
        elif image_info.sha != accelerated["revision"]:
            raise ValueError("Accelerated Flora image changed; wait for a validated release")
        image = f"hf.co/spaces/{accelerated['repo']}"
        command = (
            ["/opt/flora-venv/bin/python", "/workspace/app/flora_studio_entry.py"]
            if integrated
            else ["python", "/workspace/app/flora_speed_entry.py"]
        )
        manifest["worker_image"] = {"repo": accelerated["repo"], "revision": image_info.sha}
        environment.update(
            ANNOTATION_MANIFEST=json.dumps(manifest),
            FLORA_EXECUTION="accelerated",
            FLORA_BATCH_SIZE=str(accelerated.get("batch_size", 4)),
            FLORA_VISION_FRAME_BATCHING="1" if accelerated.get("vision_frame_batching") else "0",
            FLORA_EPISODE_PARALLELISM=str(accelerated.get("episode_parallelism", 1)),
            FLORA_GPU_BATCH_SIZE=str(
                accelerated.get("generation_batch_size", accelerated.get("batch_size", 4))
            ),
            FLORA_ENABLE_INFERENCE="1",
            FLORA_EXPECTED_CODE=json.dumps(accelerated["expected_code"]),
            FLORA_RUNTIME_SECONDS=str(
                int(timeout_minutes * 60 if timeout_minutes else plan["costs"]["hours"] * 3600)
            ),
        )
        if integrated:
            environment["FLORA_EXPECTED_KERNELS"] = json.dumps(accelerated["kernel_versions"])
            environment["FLORA_EXPECTED_RUNTIME_FILES"] = json.dumps(accelerated["runtime_files"])
    if hands:
        if hands["runtime"].get("image_source") == "studio" and hands["runtime"]["image_repo"] != space:
            raise ValueError("Integrated hand actions must use this Annotation Studio Space")
        built = api.space_info(hands["runtime"]["image_repo"])
        if built.sha != hands["runtime"]["image_revision"]:
            raise ValueError(
                "The hand image changed; wait for a pinned release and calculate a fresh estimate"
            )
        if built.runtime.stage != "RUNNING" or built.runtime.raw.get("sha") != built.sha:
            raise ValueError("The hand image is still building; wait for its checks before starting a Job")
        image = hands["runtime"]["image"]
        command = job_command(space, revision, timeout_minutes, captioner, hands=True)
    job = api.run_job(
        image=image,
        command=command,
        flavor=flavor,
        timeout=f"{timeout_minutes}m" if timeout_minutes else f"{plan['costs']['hours']}h",
        namespace=username,
        name=f"annotate-{plan['nonce'][:12]}",
        labels={
            "annotation_request": plan["nonce"],
            "app": "lerobot-annotation-studio",
            "annotation_output": output.split("/", 1)[1],
        },
        env=environment,
        secrets={"HF_TOKEN": token},
    )
    return {"id": job.id, "owner": username, "output": output}