"""Paired native-Flora/Qwen3.8 captioning on frozen ABC reference intervals.

Run in a bounded GPU Job. Input contains interval boundaries and episode tasks,
never reference caption text. Each target is saved before continuing.
"""

import base64
import io
import json
import os
import tempfile
import time
from pathlib import Path
from types import SimpleNamespace

import av
import requests
from flora_captioner import FLORA, FloraCaptioner
from huggingface_hub import HfApi, hf_hub_download
from PIL import Image

QWEN = "Qwen/Qwen3.8-27B:novita"


class NativeFrames:
    def __init__(self, entry, path):
        self.entry, self.path, self.cache = entry, path, {}

    def frames_at(self, record, timestamps):
        fps = self.entry["fps"]
        indices = [min(self.entry["length"] - 1, max(0, round(t * fps))) for t in timestamps]
        with av.open(str(self.path)) as container:
            stream = container.streams.video[0]
            stream.codec_context.thread_count = 2
            for index in sorted(set(indices) - self.cache.keys()):
                target = self.entry["video_offset"] + index / fps
                container.seek(int(target / stream.time_base), stream=stream, backward=True)
                closest, distance = None, float("inf")
                for frame in container.decode(stream):
                    if frame.pts is None:
                        continue
                    stamp = float(frame.pts * stream.time_base)
                    if abs(stamp - target) < distance:
                        closest, distance = frame, abs(stamp - target)
                    if stamp >= target:
                        break
                if closest is None or distance > 1 / fps + 1e-4:
                    raise ValueError("Source frame is missing or outside its expected clock")
                self.cache[index] = closest.to_image().convert("RGB")
        return [self.cache[index] for index in indices]


def qwen_messages(sample, record):
    content = [
        {
            "type": "text",
            "text": (
                "Write one concise, self-contained English imperative for the local operation attempted in TARGET. "
                "Identify concrete objects and the working hand. Use BEFORE/AFTER only to disambiguate TARGET; "
                "do not include their operations. Do not invent hidden objects, goals or success. "
                "Return only the instruction, on one line. All frames use one top camera. "
                f"Episode instruction: {record.episode_task}"
            ),
        }
    ]
    for images, block in zip(sample["frames"], sample["sample"]["blocks"], strict=True):
        content.append({"type": "text", "text": block["role"] + ":"})
        for frame, index in zip(images, block["original_descriptor"]["frame_indices"], strict=True):
            image = (frame if isinstance(frame, Image.Image) else Image.fromarray(frame)).convert("RGB")
            buffer = io.BytesIO()
            image.save(buffer, format="JPEG", quality=84)
            content.extend(
                [
                    {"type": "text", "text": f"Source frame {index}"},
                    {
                        "type": "image_url",
                        "image_url": {
                            "url": "data:image/jpeg;base64," + base64.b64encode(buffer.getvalue()).decode()
                        },
                    },
                ]
            )
    return [{"role": "user", "content": content}]


def qwen_caption(messages, state):
    # Reserve a full million-token context before dispatch, reconcile reported usage.
    reservation = (1_000_000 * 0.42 + 512 * 3.0) / 1e6
    if state["reserved_usd"] + reservation > 3:
        raise RuntimeError("Qwen inference cap reached")
    state["reserved_usd"] += reservation
    response = requests.post(
        "https://router.huggingface.co/v1/chat/completions",
        headers={"Authorization": "Bearer " + os.environ["HF_TOKEN"]},
        json={
            "model": QWEN,
            "messages": messages,
            "temperature": 0,
            "max_tokens": 512,
            "chat_template_kwargs": {"enable_thinking": False},
        },
        timeout=(20, 300),
    )
    if not response.ok:
        raise RuntimeError(f"Qwen HTTP {response.status_code}; no automatic paid retry")
    result = response.json()
    usage = result.get("usage", {})
    if all(type(usage.get(key)) is int for key in ("prompt_tokens", "completion_tokens")):
        cost = (usage["prompt_tokens"] * 0.42 + usage["completion_tokens"] * 3) / 1e6
        state["reserved_usd"] += cost - reservation
        state["measured_usd"] += cost
    choice = result["choices"][0]
    if choice.get("finish_reason") == "length":
        raise ValueError("Qwen response exceeded output limit")
    caption = (choice["message"].get("content") or "").strip()
    if not caption:
        raise ValueError("Qwen returned an empty caption")
    return caption


def main():
    api = HfApi()
    repo = os.environ["BENCHMARK_REPO"]
    revision = os.environ["BENCHMARK_REVISION"]
    inputs = json.loads(
        Path(hf_hub_download(repo, "inputs.json", repo_type="dataset", revision=revision)).read_text()
    )
    work = Path("/workspace/results")
    work.mkdir(parents=True, exist_ok=True)
    state = {
        "flora": FLORA,
        "qwen": QWEN,
        "input_revision": revision,
        "reserved_usd": 0.0,
        "measured_usd": 0.0,
        "episodes": [],
        "started_at": time.time(),
        "status": "running",
    }

    def save():
        (work / "run.json").write_text(json.dumps(state, indent=2))
        api.upload_folder(
            repo_id=repo,
            repo_type="dataset",
            folder_path=work,
            path_in_repo="results",
            commit_message="Checkpoint paired caption benchmark",
        )

    engine = FloraCaptioner()
    try:
        engine.load()
        for entry in inputs["episodes"]:
            path = hf_hub_download(
                inputs["repo_id"], entry["video_path"], repo_type="dataset", revision=inputs["revision"]
            )
            record = SimpleNamespace(
                episode_index=entry["episode_index"],
                episode_task=entry["instruction"],
                frame_timestamps=tuple(i / entry["fps"] for i in range(entry["length"])),
            )
            frames = NativeFrames(entry, path)
            result = {"episode_index": entry["episode_index"], "segments": [], "status": "running"}
            state["episodes"].append(result)
            with tempfile.TemporaryDirectory(prefix="flora-abc-") as temporary:
                for sample in engine.samples(record, entry["spans"], frames, entry["fps"], Path(temporary)):
                    span = entry["spans"][sample["index"]]
                    output = {**span, "index": sample["index"], "status": sample["status"]}
                    result["segments"].append(output)
                    if sample["status"] == "ready":
                        output["input_blocks"] = sample["plan"]["blocks"]
                        started = time.monotonic()
                        output["flora_caption"] = engine.infer(sample).strip()
                        output["flora_seconds"] = time.monotonic() - started
                        started = time.monotonic()
                        output["qwen_caption"] = qwen_caption(qwen_messages(sample, record), state)
                        output["qwen_seconds"] = time.monotonic() - started
                        output["status"] = "complete"
                    save()
            result["status"] = "complete"
            frames.cache.clear()
        state["status"] = "complete"
    except Exception as exc:
        state.update(status="failed", error_type=type(exc).__name__, error=str(exc)[:300])
        raise
    finally:
        state["finished_at"] = time.time()
        save()
        engine.close()


if __name__ == "__main__":
    main()