"""Independent native Flora TempSeg → Captioner workflow on a CUDA Job.

The default profile uses NVIDIA's reference decoder. The opt-in accelerated
profile batches independent requests and requires fused GDN kernels; both use
the same planner, processor, mixed precision loader and greedy decoding policy.
"""

import asyncio
import hashlib
import os
import sys
import tempfile
import threading
import time
from concurrent.futures import ThreadPoolExecutor
from fractions import Fraction
from itertools import islice
from pathlib import Path

FLORA = {
    "model_id": "nvidia/Flora-Captioner",
    "code_revision": "f736ca36925ae55d2f7915993ed1bc9f87288bb3",
    "weights_revision": "fef0bbfe11829ca7f3539d58da7a6a2c2f109dc6",
    "segmenter_id": "nvidia/Flora-TempSeg",
    "segmenter_revision": "114a1076f70eb53f4a58ac045505db3087e79347",
    "segmenter_sha256": "1d80fe544e4cd1a7ff6b0ea4be71e988972408dab81d4c8af72e3a628739d6c3",
    "segmenter_base_id": "Qwen/Qwen3.5-9B",
    "segmenter_base_revision": "c202236235762e1c871ad0ccb60c8ee5ba337b9a",
    "temporal_profile": "yam",
    "hardware": "h200",
    "resident_models": True,
    "license": "NVIDIA OneWay Noncommercial License",
    "role": "TempSeg predicts and stitches intervals; Captioner writes their labels independently",
}


def reference_runtime():
    from huggingface_hub import snapshot_download

    # Reference modules resolve __file__; Hub cache symlinks would resolve into
    # blobs/ and lose sibling config/vendor files. Materialize a real code tree.
    directory = str(
        Path(os.environ.get("FLORA_CODE_CACHE", Path.home() / ".cache/flora/code")) / FLORA["code_revision"]
    )
    snapshot_download(
        FLORA["model_id"],
        revision=FLORA["code_revision"],
        allow_patterns=["batch_utils/**"],
        local_dir=directory,
    )
    sys.path.insert(0, directory)
    from batch_utils.bootstrap import production
    from batch_utils.config import MODELS
    from batch_utils.models import Session, generate
    from batch_utils.preprocessing import instruction_suffix
    from flora_unified.unified_input import load_frames, prompt

    expected = {
        "tempseg": (FLORA["segmenter_id"], FLORA["segmenter_revision"]),
        "tempseg_base": (FLORA["segmenter_base_id"], FLORA["segmenter_base_revision"]),
        "captioner": (FLORA["model_id"], FLORA["weights_revision"]),
    }
    if any((MODELS[k]["repo"], MODELS[k]["revision"]) != value for k, value in expected.items()):
        raise ValueError("Native Flora model revisions differ from the reviewed configuration")
    if MODELS["tempseg"]["sha256"] != FLORA["segmenter_sha256"]:
        raise ValueError("Native Flora adapter checksum differs from the reviewed configuration")

    return production, Session, generate, instruction_suffix, load_frames, prompt


class TemporalPredictor:
    """Adapt LeRobot frame access to NVIDIA's unchanged window/stitch engine."""

    def __init__(self, engine, source, binding, record, frames):
        self.engine, self.source, self.binding = engine, source, binding
        self.record, self.frames = record, frames

    def predict(self, request):
        from batch_utils.preprocessing import temporal_inputs
        from flora_runtime.video_segmentation.fine_subtask import FINE_RESPONSE_SCHEMA

        if (
            request["source"]["source_id"] != self.source.source_id
            or request["model_binding"] != self.binding
        ):
            raise ValueError("Native Flora temporal request binding changed")
        production = self.engine.reference[0]
        pts_us = [Fraction(str(t)) * 1_000_000 for t in self.record.frame_timestamps]
        predictions, payloads = [], iter(request["windows"])
        while batch := list(islice(payloads, self.engine.batch_size)):
            windows, requests = [], []
            for payload in batch:
                self.engine.check_deadline()
                window = production.window_from_payload(payload, self.source)
                plan = production.temporal_frame_plan(window, pts_us)
                images = self.frames.frames_at(
                    self.record, [self.record.frame_timestamps[i] for i in plan["source_frame_indices"]]
                )
                inputs, _ = temporal_inputs(
                    images,
                    window,
                    FLORA["temporal_profile"],
                    self.engine.processor,
                    ["main"],
                    self.record.episode_task or "",
                )
                windows.append(window)
                requests.append(inputs)
            for window, prediction in zip(windows, self.engine.generate(requests, 1024), strict=True):
                predictions.append({"scope_id": window.scope_id, "prediction_json": prediction})
        return {
            "schema_version": FINE_RESPONSE_SCHEMA,
            "source_id": self.source.source_id,
            "model_io_contract": self.binding["model_io_contract"],
            "model_binding": self.binding,
            "predictions": predictions,
        }


class FloraCaptioner:
    def __init__(
        self,
        deadline=None,
        cache_models=False,
        execution="reference",
        batch_size=4,
        vision_frame_batching=False,
    ):
        if execution not in {"reference", "accelerated"}:
            raise ValueError("Unknown Flora execution profile")
        if not isinstance(batch_size, int) or isinstance(batch_size, bool) or not 1 <= batch_size <= 8:
            raise ValueError("Flora batch size must be between 1 and 8")
        self.lock = threading.RLock()
        self.deadline = deadline
        self.execution = execution
        if type(vision_frame_batching) is not bool or vision_frame_batching and execution != "accelerated":
            raise ValueError("Vision frame batching requires the accelerated Flora profile")
        self.vision_frame_batching = vision_frame_batching
        self.batch_size = batch_size if execution == "accelerated" else 1
        self.kernel_versions = {}
        if execution == "accelerated":
            from flora_acceleration import require_fast_kernels

            self.kernel_versions = require_fast_kernels()
        self.reference = reference_runtime()
        self.cache_models = cache_models or execution == "accelerated"
        self.session = None if self.cache_models else self.reference[1]()
        self.sessions = {}
        self.cached_models = {}
        self.model = self.processor = None
        self.loaded_kind = None
        self.prefetch_pool = None
        self.prefetch = None
        if execution == "accelerated":
            from batch_utils.models import snapshot, validate_manifest

            def prepare_captioner():
                path = snapshot("captioner")
                # Session.load still invokes the native validator and checks
                # every fingerprint before reusing this full checksum pass.
                validate_manifest(path)
                return path

            self.prefetch_pool = ThreadPoolExecutor(max_workers=1)
            self.prefetch = self.prefetch_pool.submit(prepare_captioner)

    def close(self):
        if self.prefetch_pool is not None:
            self.prefetch_pool.shutdown(wait=True, cancel_futures=True)
        self.model = self.processor = None
        self.cached_models.clear()
        for session in self.sessions.values():
            session.close()
        self.sessions.clear()
        if self.session is not None:
            self.session.close()
        self.loaded_kind = None

    def check_resident_capacity(self):
        import torch

        # The native models use ~18 + 51 GiB before activations. This mode
        # targets the H200, leaving room for both models and their KV caches.
        if torch.cuda.get_device_properties(0).total_memory < 120 * 1024**3:
            raise ValueError("Keeping both Flora models loaded requires an H200-class GPU with 120 GiB")

    def check_deadline(self):
        if self.deadline and time.time() >= self.deadline:
            from provider import DeadlineExceededError

            raise DeadlineExceededError("Flora Job time limit reached; completed episodes are saved.")

    def annotate(self, record, frames, fps):
        """Predict boundaries with TempSeg, then caption them without a provider call."""
        with self.lock:
            spans = self.segment(record, frames, fps)
            return self.relabel(record, spans, frames, fps)

    def segment(self, record, frames, fps):
        production = self.reference[0]
        if fps < 4 or not record.frame_timestamps:
            raise ValueError("Flora temporal segmentation needs at least 4 source FPS")
        if any(abs(t - i / fps) > 1e-4 for i, t in enumerate(record.frame_timestamps)):
            raise ValueError("The Studio Flora robot profile needs an episode-relative constant-rate clock")
        self.load("tempseg")
        source = production.AnnotationSource(
            source_id=f"studio_episode_{record.episode_index}",
            location=f"lerobot:episode/{record.episode_index}",
            # NVIDIA's scheduler accepts integral microseconds, matching its
            # native media probe. Frame counts / 30 often repeat in decimal.
            duration_s=round(len(record.frame_timestamps) / fps * 1_000_000) / 1_000_000,
            fps=float(fps),
        )
        pool = FLORA["temporal_profile"]
        binding = {
            "candidate_id": "flora_final",
            "checkpoint_sha256": FLORA["segmenter_sha256"],
            "model_io_contract": production.fine_window_profile(production.PROFILES[pool]).model_io_contract,
        }
        predictor = TemporalPredictor(self, source, binding, record, frames)
        native = production.production_segmenter(predictor, binding, pool)
        segments = asyncio.run(native.segment_async(source))
        if not segments:
            raise ValueError("Flora TempSeg did not return any subtask intervals")
        return [
            {
                "start": segment.span.start_s,
                "end": segment.span.end_s,
                "metadata": dict(segment.metadata),
            }
            for segment in segments
        ]

    def samples(self, record, spans, frames, fps, directory):
        """Native duration-adaptive BEFORE/TARGET/AFTER plans without reference text."""
        production, _, _, suffix, load_frames, prompt = self.reference
        rate = Fraction(str(fps))
        plans, dispositions = production.caption_plans(
            [
                {"start_s": span["start"], "end_s": span["end"], "metadata": span.get("metadata", {})}
                for span in spans
            ],
            {
                "fps_num": rate.numerator,
                "fps_den": rate.denominator,
                "frame_count": len(record.frame_timestamps),
            },
        )
        wanted = sorted({i for plan in plans for block in plan["blocks"] for i in block["frame_indices"]})
        images = frames.frames_at(record, [record.frame_timestamps[i] for i in wanted])
        receipts = {}
        for index, image in zip(wanted, images, strict=True):
            path = directory / f"{index:012d}.png"
            image.convert("RGB").save(path, format="PNG")
            receipts[index] = {
                "path": path.name,
                "png_sha256": hashlib.sha256(path.read_bytes()).hexdigest(),
                "size": list(image.size),
            }
        by_event = {plan["event_id"]: plan for plan in plans}
        for disposition in dispositions:
            plan = by_event.get(disposition["event_id"])
            if plan is None:
                yield {"status": disposition["status"], "index": disposition["event_index"]}
                continue
            sample = production.caption_sample(plan, receipts, ["main"])
            blocks, _ = load_frames(sample, directory)
            content = prompt(sample, self.processor.video_token) + suffix(record.episode_task)
            yield {
                "status": "ready",
                "index": plan["event_index"],
                "plan": plan,
                "sample": sample,
                "frames": blocks,
                "content": content,
            }

    def caption_inputs(self, sample):
        from transformers.video_utils import VideoMetadata

        self.check_deadline()
        proc = self.processor
        text = proc.apply_chat_template(
            [{"role": "user", "content": sample["content"]}],
            tokenize=False,
            add_generation_prompt=True,
            enable_thinking=False,
        )
        metadata = [
            VideoMetadata(
                total_num_frames=len(frames),
                fps=sample["plan"]["source_fps_num"] / sample["plan"]["source_fps_den"],
                frames_indices=list(block["block_local_timestamp_indices"]),
            )
            for frames, block in zip(sample["frames"], sample["sample"]["blocks"], strict=True)
        ]
        inputs = proc(
            text=[text],
            videos=sample["frames"],
            video_metadata=metadata,
            do_sample_frames=False,
            return_metadata=True,
            return_tensors="pt",
        )
        return inputs

    def generate(self, requests, output_limit):
        """Retain reference batch-one behavior unless acceleration is explicitly selected."""
        if self.execution == "accelerated":
            from flora_acceleration import generate_many

            return generate_many(
                self.model,
                self.processor,
                requests,
                output_limit,
                batch_size=self.batch_size,
                check_deadline=self.check_deadline,
            )
        return [self.reference[2](self.model, self.processor, inputs, output_limit) for inputs in requests]

    def infer(self, sample):
        return self.generate([self.caption_inputs(sample)], 512)[0]

    def load(self, kind="captioner"):
        self.check_deadline()
        if kind == "captioner" and self.prefetch is not None:
            self.prefetch.result()
        if self.cache_models:
            if kind not in self.cached_models:
                self.check_resident_capacity()
                # Session.load closes its own previous model. A separate native
                # Session per kind preserves both verified model instances.
                session = self.reference[1]()
                model, processor = session.load(kind)
                if self.vision_frame_batching:
                    from batch_utils.models import runtime_event
                    from flora_vision import configure_vision_batching

                    layers = configure_vision_batching(model)
                    runtime_event("vision_frames_batched", kind=kind, layers=layers, backend="sdpa")
                self.cached_models[kind] = model, processor
                self.sessions[kind] = session
            self.model, self.processor = self.cached_models[kind]
            self.loaded_kind = kind
            return
        if self.model is None or self.loaded_kind != kind:
            # Release our references before Session closes the preceding model.
            self.model = self.processor = None
            self.model, self.processor = self.session.load(kind)
            self.loaded_kind = kind

    def relabel(self, record, spans, frames, fps):
        with self.lock, tempfile.TemporaryDirectory(prefix="flora-caption-") as temporary:
            self.load()
            result = [dict(span) for span in spans]
            samples = iter(self.samples(record, spans, frames, fps, Path(temporary)))
            while batch := list(islice(samples, self.batch_size)):
                if any(sample["status"] != "ready" for sample in batch):
                    raise ValueError("Flora cannot caption a source-limited target interval")
                captions = self.generate([self.caption_inputs(sample) for sample in batch], 512)
                for sample, caption in zip(batch, captions, strict=True):
                    if not caption.strip():
                        raise ValueError("Flora returned an empty caption")
                    result[sample["index"]]["text"] = caption.strip()
            return result