"""Studio episode/media stages on the LeRobot #4892 execution contracts.

The CPU controller keeps its shared billing guard and provider semaphores in
process. Only opaque in-process handles enter operational plans; tokens and
provider objects are never serialized. Scientific identities exclude handles.
"""

import hashlib
import json
import shutil
import tarfile
import uuid
from dataclasses import asdict
from pathlib import Path

from lerobot.annotations.runtime.configs import RuntimeConfig, StageConfig
from lerobot.annotations.runtime.pipeline import run_pipeline
from lerobot.annotations.runtime.types import (
    Artifact,
    DatasetRef,
    InputItem,
    ItemResult,
    ModuleSpec,
    Outcome,
    fingerprint,
)
from lerobot.annotations.runtime.worker import accepted_items

CONTEXTS = {}


def unpack(bundle, directory):
    with tarfile.open(bundle) as archive:
        for member in archive.getmembers():
            path = Path(member.name)
            if not member.isfile() or path.is_absolute() or ".." in path.parts:
                raise ValueError("Invalid Studio checkpoint bundle")
            target = directory / path
            target.parent.mkdir(parents=True, exist_ok=True)
            target.write_bytes(archive.extractfile(member).read())


class StudioStage:
    def __init__(self, handle, phase, scientific):
        self.handle, self.phase, self.scientific = handle, phase, scientific
        self.spec = ModuleSpec("studio_" + phase, "1", "episode", {"bundle": None})

    def scientific_config(self):
        return {"phase": self.phase, **self.scientific}

    def setup(self, context):
        self.values = CONTEXTS[self.handle]

    def teardown(self):
        pass

    def process_batch(self, items, context):
        import worker

        from lerobot.annotations.steerable_pipeline.staging import EpisodeStaging

        record, row, info, root, manifest, provider, media = self.values
        results = []
        for item in items:
            directory = context.scratch / item.item_id
            directory.mkdir(parents=True, exist_ok=True)
            staging_dir = directory / "staging"
            if self.phase == "language":
                result = worker._annotate_episode(
                    record, row, info, root, manifest, provider, staging_dir, None
                )
            else:
                artifact = Artifact(**item.payload["upstream"])
                if not context.store.verify(artifact):
                    raise ValueError("Upstream checkpoint checksum mismatch")
                bundle = context.store.download(artifact, directory / "language.tar.gz")
                unpack(bundle, directory)
                bundle.unlink()
                result = json.loads((directory / "result.json").read_text())
            if self.phase == "hands":
                import math
                import time

                import av
                import pyarrow.parquet as pq
                from subset import trim_video

                from lerobot.annotations.hands.writer import attach_actions

                if worker.HAND_ANNOTATOR is None:
                    raise RuntimeError("The HF Jobs hand worker was not initialized")
                native = directory / "hand-native"
                native.mkdir(parents=True, exist_ok=True)
                video = native / "episode.mp4"
                started = time.perf_counter()
                print(
                    "HAND_RUNTIME "
                    + json.dumps(
                        {"event": "episode_video_prepare_start", "source_episode": record.episode_index}
                    ),
                    flush=True,
                )
                source_video = root / worker.video_path(row, info, manifest["camera"])
                offset = float(row[f"videos/{manifest['camera']}/from_timestamp"])
                # Standalone raw episode videos need no RGB re-encode. The
                # native writer independently verifies decoded frame count/FPS.
                with av.open(str(source_video)) as probe:
                    stream = probe.streams.video[0]
                    complete_episode = (
                        len(probe.streams.video) == 1
                        and offset == 0
                        and stream.start_time == 0
                        and stream.frames == int(row["length"])
                        and math.isclose(float(stream.average_rate or 0), float(info["fps"]), rel_tol=1e-6)
                    )
                if complete_episode:
                    video = source_video
                else:
                    trim_video(
                        source_video, video, offset, int(row["length"]), float(info["fps"]), lossless=True
                    )
                prepared = time.perf_counter()
                print(
                    "HAND_RUNTIME "
                    + json.dumps(
                        {
                            "event": "native_annotation_start",
                            "source_episode": record.episode_index,
                            "video_prepare_seconds": prepared - started,
                            "input_mode": "whole-episode-video" if complete_episode else "lossless-trim",
                        }
                    ),
                    flush=True,
                )
                try:
                    native_report = worker.HAND_ANNOTATOR.annotate(
                        video,
                        manifest["hand_actions"]["calibration"],
                        native / "work",
                        orientation=manifest["hand_actions"].get("orientation", "sensor"),
                    )
                except Exception:
                    # Native inference is offline and has no Hub credentials.
                    import logging

                    logging.exception(
                        "Native hand annotation failed for source episode %s", record.episode_index
                    )
                    raise
                inferred = time.perf_counter()
                table = pq.read_table(
                    record.data_path, filters=[("episode_index", "=", record.episode_index)]
                )
                table, summary = attach_actions(table, native_report, directory=native, fps=info["fps"])
                summary["timing_seconds"] = {
                    "source_video_prepare": prepared - started,
                    "native_annotation": inferred - prepared,
                    "native_scope": "intrinsics, rectification, camera, hands and native export; includes coordinator wait and first-use model initialization; excludes source preparation and final dataset video copying",
                }
                summary["native_components"] = {
                    key: native_report[key]
                    for key in ("camera_statistics", "hand_components")
                    if key in native_report
                }
                print(
                    "HAND_RUNTIME "
                    + json.dumps(
                        {
                            "event": "native_annotation_finished",
                            "source_episode": record.episode_index,
                            "native_seconds": inferred - prepared,
                            "valid_frames": summary["valid_frames"],
                        }
                    ),
                    flush=True,
                )
                pq.write_table(table, directory / "hand_actions.parquet")
                result["hand_actions"] = summary
                # Model assets and RGB/prediction caches never enter dataset uploads or checkpoints.
                shutil.rmtree(native)
            elif self.phase == "human_video":
                staging = EpisodeStaging(staging_dir, record.episode_index)
                frames = worker.SourceFrames(row, info, root, manifest["camera"])
                full = worker.SourceFrames(row, info, root, manifest["camera"], max_size=None)
                try:
                    entries = worker.generate_human_videos(
                        record,
                        staging,
                        provider,
                        frames,
                        full,
                        media,
                        directory / "human_videos" / f"episode_{record.episode_index:06d}",
                        manifest["human_video"]["segments_per_episode"],
                    )
                except worker.BudgetExceededError as exc:
                    entries = [
                        {
                            "source_episode_index": record.episode_index,
                            "status": "time_stopped"
                            if isinstance(exc, worker.DeadlineExceededError)
                            else "budget_stopped",
                            "files": {},
                        }
                    ]
                for entry in entries:
                    entry["files"] = {
                        kind: str(Path(path).relative_to(directory))
                        for kind, path in entry.get("files", {}).items()
                    }
                result["human_videos"] = worker.human_video_summary(entries)
                result["human_video_entries"] = entries
            (directory / "result.json").write_text(json.dumps(result))
            bundle = context.scratch / (item.item_id + ".tar.gz")
            with tarfile.open(bundle, "w:gz") as archive:
                for path in sorted(directory.rglob("*")):
                    if path.is_file():
                        archive.add(path, arcname=str(path.relative_to(directory)))
            artifact = context.write_asset(item, "bundle", bundle)
            results.append(ItemResult(item.item_id, Outcome.COMPLETED, (artifact,)))
        return results


class EpisodeSource:
    def __init__(self, record, row, info, manifest):
        self.record, self.row = record, row
        self.identity = {
            "row": row,
            "fps": info["fps"],
            "task": record.episode_task,
            "source": manifest.get("dataset"),
            "revision": manifest.get("dataset_revision"),
        }
        revision = manifest.get("dataset_revision") or fingerprint(self.identity)
        self.dataset_ref = DatasetRef(manifest.get("dataset", "studio-local"), revision)

    def discover(self, stage, store, upstream):
        payload = {}
        identity = dict(self.identity)
        if upstream:
            plan, _ = next(iter(upstream.values()))
            _, result = next(accepted_items(store, plan))
            artifact = result.artifacts[0]
            payload["upstream"] = asdict(artifact)
            identity["upstream"] = {key: value for key, value in asdict(artifact).items() if key != "path"}
        yield InputItem(str(self.record.episode_index), payload, content_digest=fingerprint(identity))


def run_episode(record, row, info, root, manifest, provider, staging_dir, media=None):
    handle = uuid.uuid4().hex
    CONTEXTS[handle] = (record, row, info, root, manifest, provider, media)
    folder = Path(__file__).parent
    scientific = {
        "options": {
            key: manifest.get(key)
            for key in (
                "features",
                "camera",
                "vqa_hz",
                "prompts",
                "model",
                "generation_options",
                "human_video",
                "captioner",
                "captioner_profile",
                "flora_execution",
                "hand_actions",
            )
        },
        "code": {
            path.name: hashlib.sha256(path.read_bytes()).hexdigest()
            for pattern in ("*.py", "*.txt", "pipeline-version.json")
            for path in sorted(folder.glob(pattern))
            if not path.name.startswith("test_")
        },
    }
    stages = [
        StageConfig(
            "language",
            "studio_runtime:StudioStage",
            {"handle": handle, "phase": "language", "scientific": scientific},
        )
    ]
    if media is not None:
        stages.append(
            StageConfig(
                "human_video",
                "studio_runtime:StudioStage",
                {"handle": handle, "phase": "human_video", "scientific": scientific},
                ("language",),
            )
        )
    if manifest.get("hand_actions"):
        stages.append(
            StageConfig(
                "hands",
                "studio_runtime:StudioStage",
                {"handle": handle, "phase": "hands", "scientific": scientific},
                (stages[-1].id,),
            )
        )
    # Existing Studio threads own concurrency; one in-process runtime worker
    # shares the same request budget. Accepted media is never automatically rebilled.
    runtime = RuntimeConfig(
        run_uri=str(staging_dir.parent / "runtime" / f"episode_{record.episode_index:06d}"),
        workers=1,
        batch_size=1,
        shard_size=1,
        max_retries=0,
    )
    try:
        store, completed = run_pipeline(EpisodeSource(record, row, info, manifest), stages, runtime)
        plan, _ = completed[stages[-1].id]
        _, result = next(accepted_items(store, plan))
        directory = staging_dir.parent / "restored" / handle
        directory.mkdir(parents=True, exist_ok=True)
        bundle = store.download(result.artifacts[0], directory / "result.tar.gz")
        unpack(bundle, directory)
        result = json.loads((directory / "result.json").read_text())
        if (directory / "hand_actions.parquet").exists():
            target = staging_dir.parent / "hand_actions" / f"episode_{record.episode_index:06d}.parquet"
            target.parent.mkdir(parents=True, exist_ok=True)
            shutil.copyfile(directory / "hand_actions.parquet", target)
        if (directory / "staging").exists():
            shutil.rmtree(staging_dir / f"episode_{record.episode_index:06d}", ignore_errors=True)
            shutil.copytree(directory / "staging", staging_dir, dirs_exist_ok=True)
        for entry in result.get("human_video_entries", []):
            for kind, relative in entry.get("files", {}).items():
                source = directory / relative
                target = staging_dir.parent / relative
                target.parent.mkdir(parents=True, exist_ok=True)
                shutil.copyfile(source, target)
                entry["files"][kind] = str(target)
        result["runtime"] = {name: summary.plan_id for name, (_, summary) in completed.items()}
        shutil.rmtree(directory)
        return result
    finally:
        CONTEXTS.pop(handle, None)