"""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 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",
}
HUMAN_VIDEOS = "Human demo videos"
COPY = "Annotated dataset copy"
EXPORT = "Annotation files only"  # Historical reports only; new runs always write Parquet.


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,
):
    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"})
    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.
    inputs = (seconds / 10 * 1500 + windows * 1500) if needs_subtasks else 0
    outputs = (seconds / 5 * 55 + windows * 120) if needs_subtasks else 0
    calls = windows if needs_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,
        "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,
):
    return hashlib.sha256(
        json.dumps(
            [
                dataset.strip(),
                start,
                count,
                sorted(features),
                camera,
                vqa_hz,
                hours,
                mode,
                model_key,
                human_video_segments,
                human_video_parallel,
            ],
            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,
):
    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.")
    hardware = next(h for h in HfApi(token=token or False).list_jobs_hardware() if h.name == "cpu-upgrade")
    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,
    )
    return {
        "selection": selection,
        "features": list(features),
        "camera": camera,
        "vqa_hz": vqa_hz,
        "mode": mode,
        "prompts": prompts,
        "costs": costs,
        "model_profile": profile,
        "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,
        ),
    }


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):
    # 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','pipeline-*'])"
    )
    return [
        "bash",
        "-lc",
        "\n".join(
            [
                "set -euo pipefail",
                "export HF_HUB_DISABLE_XET=1 HF_HUB_DISABLE_PROGRESS_BARS=1",
                f"export ANNOTATION_DEADLINE_EPOCH=$(( $(date +%s) + {timeout_minutes * 60 - 30} ))"
                if timeout_minutes
                else ":",
                "mkdir -p /workspace/app /workspace/source",
                "cd /workspace",
                "pip install --quiet huggingface_hub==1.26.0",
                "python -c " + shlex.quote(download),
                "pip install --quiet torch==2.11.0 torchvision==0.26.0 --index-url https://download.pytorch.org/whl/cpu",
                "pip install --quiet -r /workspace/app/worker-requirements.txt",
                "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.")
    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 CPU 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.")
    hardware = next(h for h in api.list_jobs_hardware() if h.name == "cpu-upgrade")
    if hardware.unit_label != "minute" or hardware.unit_cost_usd > plan["costs"]["cpu_rate"]:
        raise ValueError("CPU 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.")
    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
    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 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": profile["route"],
            "model_profile": profile,
            "generation_options": profile["generation_options"],
            "request_concurrency": profile["concurrency"],
            "episode_parallelism": profile["concurrency"],
            "space": space,
            "space_revision": revision,
            "owner": username,
            "runtime_limit_minutes": timeout_minutes,
        }
    )
    job = api.run_job(
        image="python:3.12-slim",
        command=job_command(space, revision, timeout_minutes),
        flavor="cpu-upgrade",
        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={"ANNOTATION_MANIFEST": json.dumps(manifest)},
        secrets={"HF_TOKEN": token},
    )
    return {"id": job.id, "owner": username, "output": output}