"""Captioners: Florence-2 for images, Qwen3.5 vision-language models for whole videos."""

import re

import numpy as np
import torch
from PIL import Image
from torchcodec.decoders import VideoDecoder
from transformers import AutoProcessor, Florence2ForConditionalGeneration, Qwen3_5ForConditionalGeneration

from hub import model_path

FLORENCE = {"dir": "florence-2-large", "repo": "florence-community/Florence-2-large"}
CAPTION_TASKS = {"Brief": "<CAPTION>", "Detailed": "<DETAILED_CAPTION>", "More detailed": "<MORE_DETAILED_CAPTION>"}
DEFAULT_CAPTION_DETAIL = "Brief"
CAPTION_MAX_NEW_TOKENS = 256

# All Qwen3.5 vision-language models, loaded from the Hub cache (too large to vendor).
# Marlin-2B is a Qwen3.5-2B fine-tune trained on one fixed prompt; it answers with a
# Scene paragraph and <start - end> events, which format_marlin turns into Markdown.
VIDEO_CAPTIONERS = {
    "Qwen3.5-4B": {"repo": "Qwen/Qwen3.5-4B", "marlin": False},
    "Qwen3.5-9B": {"repo": "Qwen/Qwen3.5-9B", "marlin": False},
    "Marlin-2B (scene + events)": {"repo": "NemoStation/Marlin-2B", "marlin": True},
}
DEFAULT_VIDEO_CAPTIONERS = ["Qwen3.5-4B"]
DEFAULT_VIDEO_PROMPT = "Write a one-sentence caption that describes this video."
MARLIN_PROMPT = (
    "Provide a spatial description of this clip followed by time-ranged events.\n"
    "For each event, give the time range as <start - end> and a short description."
)
MARLIN_PIXELS_PER_FRAME = 200704  # ~448x448, the training-time cap
VIDEO_CAPTION_MAX_NEW_TOKENS = {"qwen": 256, "marlin": 1024}
VIDEO_MIN_FRAMES = 4

_CACHE: dict[str, tuple] = {}

_THINK_BLOCK = re.compile(r"^\s*<think>.*?</think>\s*", re.DOTALL)
_THINK_TAG = re.compile(r"^\s*<think>\s*")
_MARLIN_EVENT = re.compile(
    r"^\s*<?\s*(\d+\.?\d*)\s*(?:seconds?|secs?|s)?\s*-\s*(\d+\.?\d*)\s*(?:seconds?|secs?|s)?\s*>?\s*[:\-]?\s*(.+?)\s*$"
)


def needs_prompt(captioners: list[str]) -> bool:
    """True when any ticked captioner takes the user's prompt (Marlin has its own)."""
    return any(not VIDEO_CAPTIONERS[n]["marlin"] for n in captioners or [])


def load_image_captioner():
    """Load the Florence-2 (processor, model) once; kept on CPU until a GPU call."""
    folder = FLORENCE["dir"]
    if folder not in _CACHE:
        path = model_path(FLORENCE)
        _CACHE[folder] = (AutoProcessor.from_pretrained(path), Florence2ForConditionalGeneration.from_pretrained(path).eval())
    return _CACHE[folder]


def load_video_captioner(repo: str):
    """Load a Qwen3.5 video captioner (processor, model) once in bf16; kept on CPU until a GPU call."""
    if repo not in _CACHE:
        processor = AutoProcessor.from_pretrained(repo)
        model = Qwen3_5ForConditionalGeneration.from_pretrained(repo, dtype=torch.bfloat16).eval()
        _CACHE[repo] = (processor, model)
    return _CACHE[repo]


@torch.inference_mode()
def caption_image(processor, model, rgb: np.ndarray, detail: str, device: str) -> str:
    """Caption one RGB frame with Florence-2 at the given detail level."""
    task = CAPTION_TASKS[detail]
    pil = Image.fromarray(rgb)
    inputs = processor(text=task, images=pil, return_tensors="pt").to(device)
    generated = model.generate(**inputs, max_new_tokens=CAPTION_MAX_NEW_TOKENS, num_beams=3)
    text = processor.batch_decode(generated, skip_special_tokens=False)[0]
    parsed = processor.post_process_generation(text, task=task, image_size=pil.size)
    return parsed[task].strip()


def count_frames(path: str, fps: float, max_frames: int) -> int:
    """Frames to sample at `fps`, clamped to [VIDEO_MIN_FRAMES, max_frames] and the video length.

    The Qwen3.5 video processor ignores per-call min/max frame limits, so the clamp
    happens here and the result goes to the processor as `num_frames` with `fps` cleared.
    """
    metadata = VideoDecoder(path).metadata
    wanted = int(metadata.duration_seconds * fps)
    return min(max(wanted, VIDEO_MIN_FRAMES), int(max_frames), metadata.num_frames)


@torch.inference_mode()
def caption_video(processor, model, path: str, prompt: str, num_frames: int, device: str, marlin: bool) -> str:
    """Caption one video with a Qwen3.5 model from `num_frames` evenly sampled frames."""
    messages = [{"role": "user", "content": [{"type": "video", "path": path}, {"type": "text", "text": prompt}]}]
    processor_kwargs = {"num_frames": num_frames, "fps": None, "cap_pixels_per_frame": True}
    template_kwargs = {}
    if marlin:
        processor_kwargs["size"] = {"shortest_edge": 4096, "longest_edge": MARLIN_PIXELS_PER_FRAME * num_frames}
    else:
        template_kwargs["enable_thinking"] = False
    inputs = processor.apply_chat_template(
        messages,
        add_generation_prompt=True,
        tokenize=True,
        return_dict=True,
        return_tensors="pt",
        processor_kwargs=processor_kwargs,
        **template_kwargs,
    ).to(device)
    max_new_tokens = VIDEO_CAPTION_MAX_NEW_TOKENS["marlin" if marlin else "qwen"]
    output_ids = model.generate(**inputs, max_new_tokens=max_new_tokens, do_sample=False)
    text = processor.decode(output_ids[0, inputs["input_ids"].shape[1]:], skip_special_tokens=True)
    return strip_thinking(text)


def strip_thinking(text: str) -> str:
    """Drop a leading `<think>...</think>` block, or a bare `<think>` prefix, as Marlin emits."""
    return _THINK_TAG.sub("", _THINK_BLOCK.sub("", text)).strip()


def format_marlin(text: str) -> str:
    """Turn Marlin's `Scene: ... Events: <a - b> ...` output into Markdown."""
    lines = []
    for raw in text.splitlines():
        line = raw.strip()
        if not line:
            continue
        event = _MARLIN_EVENT.match(line)
        if event:
            start, end, desc = float(event.group(1)), float(event.group(2)), event.group(3)
            lines.append(f"- `{start:5.1f}s - {end:5.1f}s` {desc}")
        elif line.lower().startswith(("scene:", "events:")):
            head, _, rest = line.partition(":")
            lines.append(f"**{head.strip()}**" + (f" {rest.strip()}" if rest.strip() else ""))
        else:
            lines.append(line)
    return "\n".join(lines)