from __future__ import annotations

import json
import sys
from pathlib import Path
from typing import Iterable

import numpy as np
import torch
import torch.nn.functional as F
from PIL import Image
from safetensors import safe_open
from transformers import AutoProcessor, AutoTokenizer, WhisperFeatureExtractor

IMAGE_SUFFIXES = {".jpg", ".jpeg", ".png", ".bmp", ".webp", ".gif", ".tif", ".tiff"}
VIDEO_SUFFIXES = {".mp4", ".mov", ".avi", ".mkv", ".webm", ".flv", ".wmv"}
AUDIO_SUFFIXES = {".wav", ".mp3", ".flac", ".ogg", ".m4a", ".opus", ".aac"}
TASK_NAMES = ("retrieval", "text-matching", "clustering", "classification")
PROMPT_PREFIX = {
    "query": "Query: ",
    "document": "Document: ",
}


def repo_root() -> Path:
    return Path(__file__).resolve().parent.parent


def resolve_device(device: str | None = None) -> torch.device:
    if device:
        return torch.device(device)
    if torch.cuda.is_available():
        return torch.device("cuda")
    return torch.device("cpu")


def torch_dtype_from_name(name: str | None) -> torch.dtype:
    if not name:
        return torch.bfloat16
    lowered = name.lower()
    mapping = {
        "bf16": torch.bfloat16,
        "bfloat16": torch.bfloat16,
        "fp16": torch.float16,
        "float16": torch.float16,
        "fp32": torch.float32,
        "float32": torch.float32,
    }
    if lowered not in mapping:
        raise ValueError(f"Unsupported dtype: {name}")
    return mapping[lowered]


def _ensure_model_repo_on_sys_path(model_dir: Path) -> None:
    model_dir = model_dir.resolve()
    model_dir_str = str(model_dir)
    if model_dir_str not in sys.path:
        sys.path.insert(0, model_dir_str)


def load_model_classes(model_dir: str | Path):
    model_dir = Path(model_dir)
    _ensure_model_repo_on_sys_path(model_dir)
    from modeling_jina_embeddings_v5_omni import (
        JinaEmbeddingsV5OmniBase,
        JinaEmbeddingsV5OmniConfig,
    )

    return JinaEmbeddingsV5OmniConfig, JinaEmbeddingsV5OmniBase


def merge_task_lora_inplace(model: torch.nn.Module, model_dir: str | Path, task: str) -> None:
    model_dir = Path(model_dir)
    if task not in TASK_NAMES:
        raise ValueError(f"Unsupported task: {task}")

    adapter_dir = model_dir / "adapters" / task
    with open(adapter_dir / "adapter_config.json", "r", encoding="utf-8") as handle:
        adapter_cfg = json.load(handle)

    scale = adapter_cfg["lora_alpha"] / adapter_cfg["r"]
    with safe_open(str(adapter_dir / "adapter_model.safetensors"), framework="pt") as handle:
        adapter = {key: handle.get_tensor(key) for key in handle.keys()}

    with torch.no_grad():
        for a_key in list(adapter.keys()):
            if "lora_A" not in a_key:
                continue
            b_key = a_key.replace("lora_A", "lora_B")
            if b_key not in adapter:
                continue
            target_name = a_key.replace("base_model.model.", "").replace(".lora_A.weight", ".weight")
            parts = target_name.split(".")
            module = model
            for part in parts[:-2]:
                module = module[int(part)] if part.isdigit() else getattr(module, part)
            proj = getattr(module, parts[-2])
            delta = (adapter[b_key].float() @ adapter[a_key].float()) * scale
            proj.weight.data.add_(delta.to(device=proj.weight.device, dtype=proj.weight.dtype))


def load_base_model(
    model_dir: str | Path,
    modality: str,
    task: str,
    dtype_name: str = "bfloat16",
    device: str | None = None,
    merge_task_lora: bool = True,
    attn_implementation: str | None = None,
):
    model_dir = Path(model_dir)
    config_cls, base_cls = load_model_classes(model_dir)
    config = config_cls.from_pretrained(str(model_dir))
    config.modality = modality
    if attn_implementation is not None:
        config._attn_implementation = attn_implementation
        for attr_name in ("vision_config", "text_config", "audio_config"):
            sub_cfg = getattr(config, attr_name, None)
            if sub_cfg is not None:
                setattr(sub_cfg, "_attn_implementation", attn_implementation)
    dtype = torch_dtype_from_name(dtype_name)
    model = base_cls.from_pretrained(str(model_dir), config=config, torch_dtype=dtype)
    if merge_task_lora:
        merge_task_lora_inplace(model, model_dir, task)
    model.set_task(task)
    resolved_device = resolve_device(device)
    model.to(device=resolved_device)
    if dtype == torch.float32:
        model.float()
    else:
        model.to(dtype=dtype)
    model.eval()
    return model


def load_processor(model_dir: str | Path, pixel_budget: int | None = None):
    kwargs = {"trust_remote_code": True}
    if pixel_budget is not None:
        kwargs["min_pixels"] = pixel_budget
        kwargs["max_pixels"] = pixel_budget
    return AutoProcessor.from_pretrained(str(model_dir), **kwargs)


def load_tokenizer(model_dir: str | Path):
    return AutoTokenizer.from_pretrained(str(model_dir), trust_remote_code=True)


def prompt_prefix(prompt_name: str) -> str:
    if prompt_name not in PROMPT_PREFIX:
        raise ValueError(f"Unsupported prompt name: {prompt_name}")
    return PROMPT_PREFIX[prompt_name]


def get_image_token_id(config) -> int:
    token_id = getattr(config, "image_token_id", None)
    if token_id is None:
        token_id = getattr(config, "image_token_index", None)
    if token_id is None:
        raise AttributeError("Config does not define image_token_id or image_token_index")
    return int(token_id)


def get_video_token_id(config) -> int:
    token_id = getattr(config, "video_token_id", None)
    if token_id is not None:
        return int(token_id)
    return get_image_token_id(config)


def build_image_prompt(processor, prompt_name: str = "query") -> str:
    image_token = getattr(processor, "image_token", "<|image_pad|>")
    text = f"{prompt_prefix(prompt_name)}<|vision_start|>{image_token}<|vision_end|>"
    return processor.apply_chat_template(
        [{"role": "user", "content": text}],
        tokenize=False,
        add_generation_prompt=False,
    )


def build_video_prompt(processor, prompt_name: str = "query") -> str:
    video_token = getattr(processor, "video_token", None) or getattr(processor, "image_token", "<|image_pad|>")
    text = f"{prompt_prefix(prompt_name)}<|vision_start|>{video_token}<|vision_end|>"
    return processor.apply_chat_template(
        [{"role": "user", "content": text}],
        tokenize=False,
        add_generation_prompt=False,
    )


def resize_image_object(image: Image.Image, width: int, height: int) -> Image.Image:
    resampling = getattr(Image, "Resampling", Image).BICUBIC
    return image.convert("RGB").resize((width, height), resampling)


def open_and_resize_image(image_path: str | Path, width: int, height: int) -> Image.Image:
    image = Image.open(image_path)
    return resize_image_object(image, width=width, height=height)


def is_video_path(path: str | Path) -> bool:
    return Path(path).suffix.lower() in VIDEO_SUFFIXES


def is_audio_path(path: str | Path) -> bool:
    return Path(path).suffix.lower() in AUDIO_SUFFIXES


def prepare_fixed_image_inputs(
    model_dir: str | Path,
    image_path: str | Path,
    width: int,
    height: int,
    prompt_name: str = "query",
    processor=None,
):
    pixel_budget = width * height
    processor = processor or load_processor(model_dir, pixel_budget=pixel_budget)
    prompt = build_image_prompt(processor, prompt_name=prompt_name)
    image = open_and_resize_image(image_path, width=width, height=height)
    inputs = processor(images=image, text=prompt, return_tensors="pt", truncation=False)
    return inputs, {"pixel_budget": pixel_budget, "prompt": prompt}


def prepare_text_inputs(
    model_dir: str | Path,
    text: str,
    prompt_name: str = "query",
    max_length: int = 32768,
):
    tokenizer = load_tokenizer(model_dir)
    encoded = tokenizer(
        [f"{prompt_prefix(prompt_name)}{text}"],
        return_tensors="pt",
        padding=True,
        truncation=True,
        max_length=max_length,
    )
    return encoded


def _extract_audio_from_video(video_path: str | Path, target_sr: int = 16000) -> np.ndarray:
    import av
    from av.audio.resampler import AudioResampler

    container = av.open(str(video_path))
    resampler = AudioResampler(format="flt", layout="mono", rate=target_sr)
    samples: list[np.ndarray] = []
    try:
        for frame in container.decode(audio=0):
            for resampled in resampler.resample(frame):
                samples.append(resampled.to_ndarray().flatten())
        for resampled in resampler.resample(None):
            samples.append(resampled.to_ndarray().flatten())
    finally:
        container.close()
    if not samples:
        return np.zeros((0,), dtype=np.float32)
    return np.concatenate(samples).astype(np.float32)


def load_audio_array(audio_path: str | Path, target_sr: int = 16000) -> np.ndarray:
    audio_path = Path(audio_path)
    if is_video_path(audio_path):
        audio = _extract_audio_from_video(audio_path, target_sr=target_sr)
    else:
        import librosa

        audio, _ = librosa.load(str(audio_path), sr=target_sr, mono=True)
        audio = audio.astype(np.float32)
    if audio.size == 0:
        raise ValueError(f"No audio samples decoded from: {audio_path}")
    peak = float(np.max(np.abs(audio)))
    if peak > 1.0:
        audio = audio / peak
    return audio.astype(np.float32)


def extract_uniform_video_frames(
    video_path: str | Path,
    *,
    num_frames: int,
    width: int,
    height: int,
) -> tuple[list[Image.Image], list[int], int]:
    if num_frames <= 0:
        raise ValueError(f"num_frames must be positive, got {num_frames}")

    video_path = Path(video_path)
    if video_path.is_dir():
        frame_paths = sorted(
            path for path in video_path.iterdir()
            if path.is_file() and path.suffix.lower() in IMAGE_SUFFIXES
        )
        frames = [Image.open(path).convert("RGB") for path in frame_paths]
    else:
        import av

        container = av.open(str(video_path))
        try:
            frames = [frame.to_image().convert("RGB") for frame in container.decode(video=0)]
        finally:
            container.close()

    if not frames:
        raise ValueError(f"No video frames decoded from: {video_path}")

    if len(frames) <= num_frames:
        indices = list(range(len(frames)))
    else:
        indices = np.linspace(0, len(frames) - 1, num_frames, dtype=int).tolist()
    selected = [resize_image_object(frames[index], width=width, height=height) for index in indices]
    return selected, indices, len(frames)


def expand_multimodal_placeholder_ids(
    tokenizer,
    prompt: str,
    *,
    placeholder_token_id: int,
    placeholder_count: int,
) -> list[int]:
    encoded = tokenizer([prompt], return_tensors="pt", padding=False, truncation=False)
    base_ids = encoded["input_ids"][0].tolist()
    current_count = sum(1 for token_id in base_ids if token_id == placeholder_token_id)
    if current_count == placeholder_count:
        return base_ids
    if current_count != 1:
        raise ValueError(
            f"Unexpected placeholder token count in prompt: got {current_count}, expected 1 or {placeholder_count}"
        )

    expanded: list[int] = []
    replaced = False
    for token_id in base_ids:
        if token_id == placeholder_token_id and not replaced:
            expanded.extend([placeholder_token_id] * placeholder_count)
            replaced = True
            continue
        expanded.append(token_id)
    return expanded


def build_audio_prompt(processor, tokenizer, config, prompt_name: str = "query", n_audio_tokens: int = 750) -> str:
    start = tokenizer.convert_ids_to_tokens(config.audio_start_token_id)
    token = tokenizer.convert_ids_to_tokens(config.audio_token_id)
    end = tokenizer.convert_ids_to_tokens(config.audio_end_token_id)
    text = f"{prompt_prefix(prompt_name)}{start}{token * n_audio_tokens}{end}"
    return processor.apply_chat_template(
        [{"role": "user", "content": text}],
        tokenize=False,
        add_generation_prompt=False,
    )


def prepare_fixed_audio_inputs(
    model_dir: str | Path,
    audio_path: str | Path,
    prompt_name: str = "query",
    max_frames: int = 3000,
    processor=None,
    tokenizer=None,
    config=None,
):
    model_dir = Path(model_dir)
    processor = processor or load_processor(model_dir)
    tokenizer = tokenizer or load_tokenizer(model_dir)
    if config is None:
        config_cls, _ = load_model_classes(model_dir)
        config = config_cls.from_pretrained(str(model_dir))

    audio = load_audio_array(audio_path)
    feat_ext = WhisperFeatureExtractor(feature_size=128)
    max_samples = int(max_frames) * int(feat_ext.hop_length)
    audio_inputs = feat_ext(
        audio,
        sampling_rate=16000,
        return_tensors="pt",
        padding="max_length",
        max_length=max_samples,
        return_attention_mask=True,
    )
    input_features = audio_inputs["input_features"]
    feature_attention_mask = audio_inputs["attention_mask"]
    if input_features.shape[-1] != max_frames:
        raise ValueError(f"Expected {max_frames} mel frames, got {input_features.shape[-1]}")

    n_audio_tokens = input_features.shape[-1] // 4
    prompt = build_audio_prompt(processor, tokenizer, config, prompt_name=prompt_name, n_audio_tokens=n_audio_tokens)
    text_inputs = processor(text=[prompt], return_tensors="pt", padding=False, truncation=False)
    text_inputs["input_features"] = input_features
    text_inputs["feature_attention_mask"] = feature_attention_mask
    return text_inputs, {
        "prompt": prompt,
        "n_audio_tokens": n_audio_tokens,
        "feature_frames": int(input_features.shape[-1]),
        "real_frames": int(feature_attention_mask.sum().item()),
    }


def prepare_fixed_video_inputs(
    model_dir: str | Path,
    video_path: str | Path,
    *,
    num_frames: int,
    width: int,
    height: int,
    prompt_name: str = "query",
    max_prefill_tokens: int | None = None,
    processor=None,
    tokenizer=None,
    config=None,
):
    model_dir = Path(model_dir)
    processor = processor or load_processor(model_dir, pixel_budget=width * height)
    tokenizer = tokenizer or load_tokenizer(model_dir)
    if config is None:
        config_cls, _ = load_model_classes(model_dir)
        config = config_cls.from_pretrained(str(model_dir))

    frames, sampled_indices, total_frames = extract_uniform_video_frames(
        video_path,
        num_frames=num_frames,
        width=width,
        height=height,
    )
    image_prompt = build_image_prompt(processor, prompt_name=prompt_name)
    video_prompt = build_video_prompt(processor, prompt_name=prompt_name)

    per_frame_pixel_values: list[torch.Tensor] = []
    frame_grid = None
    for frame in frames:
        frame_inputs = processor(images=frame, text=image_prompt, return_tensors="pt", truncation=False)
        pixel_values = frame_inputs["pixel_values"].cpu()
        image_grid_thw = frame_inputs["image_grid_thw"].cpu()
        if frame_grid is None:
            frame_grid = image_grid_thw[0]
        elif not torch.equal(frame_grid, image_grid_thw[0]):
            raise ValueError("All selected video frames must produce the same image_grid_thw for static-shape export")
        per_frame_pixel_values.append(pixel_values)

    if frame_grid is None or not per_frame_pixel_values:
        raise ValueError(f"No usable video frames after preprocessing: {video_path}")

    spatial_merge_size = int(getattr(getattr(config, "vision_config", None), "spatial_merge_size", 2) or 2)
    tokens_per_frame = int(frame_grid.prod().item()) // (spatial_merge_size ** 2)
    video_token_id = get_video_token_id(config)
    base_ids = tokenizer([video_prompt], return_tensors="pt", padding=False, truncation=False)["input_ids"][0].tolist()
    base_placeholder_count = sum(1 for token_id in base_ids if token_id == video_token_id)
    prompt_overhead = len(base_ids) - base_placeholder_count

    requested_num_frames = len(per_frame_pixel_values)
    max_frames_by_budget = None
    if max_prefill_tokens is not None:
        max_frames_by_budget = max(1, (int(max_prefill_tokens) - prompt_overhead) // tokens_per_frame)
        if requested_num_frames > max_frames_by_budget:
            per_frame_pixel_values = per_frame_pixel_values[:max_frames_by_budget]
            sampled_indices = sampled_indices[:max_frames_by_budget]

    used_num_frames = len(per_frame_pixel_values)
    total_video_tokens = used_num_frames * tokens_per_frame
    input_ids = torch.tensor(
        [
            expand_multimodal_placeholder_ids(
                tokenizer,
                video_prompt,
                placeholder_token_id=video_token_id,
                placeholder_count=total_video_tokens,
            )
        ],
        dtype=torch.long,
    )
    attention_mask = torch.ones_like(input_ids)
    pixel_values_frames = torch.stack([value for value in per_frame_pixel_values], dim=0)
    pixel_values_videos = torch.cat([value for value in per_frame_pixel_values], dim=0)
    video_grid_thw = torch.tensor(
        [[used_num_frames, int(frame_grid[1].item()), int(frame_grid[2].item())]],
        dtype=frame_grid.dtype,
    )

    return {
        "input_ids": input_ids,
        "attention_mask": attention_mask,
        "image_grid_thw": frame_grid.unsqueeze(0),
        "pixel_values_frames": pixel_values_frames,
        "pixel_values_videos": pixel_values_videos,
        "video_grid_thw": video_grid_thw,
    }, {
        "prompt": video_prompt,
        "requested_num_frames": int(num_frames),
        "used_num_frames": used_num_frames,
        "sampled_frame_indices": sampled_indices,
        "decoded_total_frames": int(total_frames),
        "tokens_per_frame": tokens_per_frame,
        "total_video_tokens": total_video_tokens,
        "prompt_overhead": prompt_overhead,
        "max_prefill_tokens": int(max_prefill_tokens) if max_prefill_tokens is not None else None,
        "max_frames_by_budget": int(max_frames_by_budget) if max_frames_by_budget is not None else None,
    }


def move_tensor_inputs(inputs: dict, device: torch.device) -> dict:
    return {key: value.to(device) if torch.is_tensor(value) else value for key, value in inputs.items()}


def last_token_embedding(last_hidden_state: torch.Tensor, attention_mask: torch.Tensor, truncate_dim: int | None = None):
    seq_lens = attention_mask.sum(dim=1) - 1
    pooled = last_hidden_state[torch.arange(last_hidden_state.shape[0], device=last_hidden_state.device), seq_lens]
    if truncate_dim is not None:
        pooled = pooled[:, :truncate_dim]
    return F.normalize(pooled, dim=-1)


def build_bidi_attention_mask(
    inputs_embeds: torch.Tensor,
    attention_mask: torch.Tensor | None,
) -> torch.Tensor | None:
    if attention_mask is None or attention_mask.dim() != 2:
        return attention_mask
    dtype = inputs_embeds.dtype
    seq_len = inputs_embeds.shape[1]
    bidi = attention_mask[:, None, None, :].to(dtype=dtype)
    bidi = (1.0 - bidi) * torch.finfo(dtype).min
    return bidi.expand(-1, -1, seq_len, -1)


def run_embedding_with_replaced_multimodal_tokens(
    model: torch.nn.Module,
    inputs: dict,
    replacement_tokens: np.ndarray | torch.Tensor,
    modality: str,
    truncate_dim: int | None = None,
) -> torch.Tensor:
    if modality not in {"vision", "audio", "video"}:
        raise ValueError(f"Unsupported modality: {modality}")

    device = next(model.parameters()).device
    torch_inputs = move_tensor_inputs(inputs, device)
    input_ids = torch_inputs["input_ids"]
    attention_mask = torch_inputs.get("attention_mask")
    inputs_embeds = model.get_input_embeddings()(input_ids)

    if isinstance(replacement_tokens, np.ndarray):
        replacement_tokens = torch.from_numpy(replacement_tokens)
    replacement_tokens = replacement_tokens.to(device=device, dtype=inputs_embeds.dtype)
    flat_tokens = replacement_tokens.reshape(-1, replacement_tokens.shape[-1])

    if modality == "vision":
        token_id = get_image_token_id(model.config)
    elif modality == "audio":
        token_id = model.config.audio_token_id
    else:
        token_id = get_video_token_id(model.config)
    token_count = int((input_ids == token_id).sum().item())
    if token_count != flat_tokens.shape[0]:
        raise ValueError(
            f"Replacement token count mismatch for {modality}: "
            f"mask expects {token_count}, got {flat_tokens.shape[0]}"
        )

    mask = (input_ids == token_id).unsqueeze(-1).expand_as(inputs_embeds)
    inputs_embeds = inputs_embeds.masked_scatter(mask, flat_tokens)
    language_model_attention_mask = build_bidi_attention_mask(inputs_embeds, attention_mask)

    with torch.no_grad():
        out = model.language_model(
            inputs_embeds=inputs_embeds,
            attention_mask=language_model_attention_mask,
        )
    return last_token_embedding(out[0], attention_mask, truncate_dim=truncate_dim)


def run_text_embedding(
    model_dir: str | Path,
    text: str,
    task: str = "retrieval",
    prompt_name: str = "query",
    dtype_name: str = "bfloat16",
    device: str | None = None,
    truncate_dim: int | None = None,
):
    model = load_base_model(model_dir, modality="text", task=task, dtype_name=dtype_name, device=device)
    inputs = prepare_text_inputs(model_dir, text=text, prompt_name=prompt_name)
    inputs = move_tensor_inputs(inputs, next(model.parameters()).device)
    with torch.no_grad():
        outputs = model(**inputs)
    return last_token_embedding(outputs.last_hidden_state, inputs["attention_mask"], truncate_dim=truncate_dim)


def run_image_embedding(
    model_dir: str | Path,
    image_path: str | Path,
    width: int,
    height: int,
    task: str = "retrieval",
    prompt_name: str = "query",
    dtype_name: str = "bfloat16",
    device: str | None = None,
    truncate_dim: int | None = None,
):
    model = load_base_model(model_dir, modality="vision", task=task, dtype_name=dtype_name, device=device)
    inputs, _ = prepare_fixed_image_inputs(model_dir, image_path=image_path, width=width, height=height, prompt_name=prompt_name)
    inputs = move_tensor_inputs(inputs, next(model.parameters()).device)
    with torch.no_grad():
        outputs = model(**inputs)
    return last_token_embedding(outputs.last_hidden_state, inputs["attention_mask"], truncate_dim=truncate_dim)


def run_audio_embedding(
    model_dir: str | Path,
    audio_path: str | Path,
    task: str = "retrieval",
    prompt_name: str = "query",
    dtype_name: str = "bfloat16",
    device: str | None = None,
    truncate_dim: int | None = None,
    max_frames: int = 3000,
):
    model_dir = Path(model_dir)
    model = load_base_model(model_dir, modality="audio", task=task, dtype_name=dtype_name, device=device)
    processor = load_processor(model_dir)
    tokenizer = load_tokenizer(model_dir)
    inputs, _ = prepare_fixed_audio_inputs(
        model_dir,
        audio_path=audio_path,
        prompt_name=prompt_name,
        max_frames=max_frames,
        processor=processor,
        tokenizer=tokenizer,
        config=model.config,
    )
    inputs = move_tensor_inputs(inputs, next(model.parameters()).device)
    del inputs["feature_attention_mask"]
    with torch.no_grad():
        outputs = model(**inputs)
    return last_token_embedding(outputs.last_hidden_state, inputs["attention_mask"], truncate_dim=truncate_dim)


def get_video_tokens_from_processed(model, pixel_values_videos: torch.Tensor, video_grid_thw: torch.Tensor) -> torch.Tensor:
    with torch.no_grad():
        features = model.get_image_features(pixel_values_videos, video_grid_thw)
    if len(features) != 1:
        raise ValueError(f"Expected one video sample, got {len(features)}")
    return features[0].unsqueeze(0)


def run_video_embedding(
    model_dir: str | Path,
    video_path: str | Path,
    *,
    num_frames: int,
    width: int,
    height: int,
    task: str = "retrieval",
    prompt_name: str = "query",
    dtype_name: str = "bfloat16",
    device: str | None = None,
    truncate_dim: int | None = None,
    max_prefill_tokens: int | None = None,
):
    model_dir = Path(model_dir)
    model = load_base_model(model_dir, modality="vision", task=task, dtype_name=dtype_name, device=device)
    processor = load_processor(model_dir, pixel_budget=width * height)
    tokenizer = load_tokenizer(model_dir)
    inputs, _ = prepare_fixed_video_inputs(
        model_dir,
        video_path=video_path,
        num_frames=num_frames,
        width=width,
        height=height,
        prompt_name=prompt_name,
        max_prefill_tokens=max_prefill_tokens,
        processor=processor,
        tokenizer=tokenizer,
        config=model.config,
    )
    inputs = move_tensor_inputs(inputs, next(model.parameters()).device)
    replacement_tokens = get_video_tokens_from_processed(
        model,
        inputs["pixel_values_videos"],
        inputs["video_grid_thw"],
    )
    lm_inputs = {
        "input_ids": inputs["input_ids"],
        "attention_mask": inputs["attention_mask"],
    }
    return run_embedding_with_replaced_multimodal_tokens(
        model,
        lm_inputs,
        replacement_tokens,
        modality="video",
        truncate_dim=truncate_dim,
    )


def get_vision_tokens_from_processed(model, pixel_values: torch.Tensor, image_grid_thw: torch.Tensor) -> torch.Tensor:
    with torch.no_grad():
        features = model.get_image_features(pixel_values, image_grid_thw)
    if len(features) != 1:
        raise ValueError(f"Expected one image, got {len(features)}")
    return features[0].unsqueeze(0)


def to_numpy(tensor: torch.Tensor) -> np.ndarray:
    return tensor.detach().cpu().float().numpy()


def list_image_files(dataset_dir: str | Path) -> list[Path]:
    dataset_dir = Path(dataset_dir)
    return sorted(
        path for path in dataset_dir.rglob("*")
        if path.is_file() and path.suffix.lower() in IMAGE_SUFFIXES
    )


def cosine_similarity(a: np.ndarray, b: np.ndarray) -> float:
    lhs = a.reshape(-1).astype(np.float64)
    rhs = b.reshape(-1).astype(np.float64)
    denom = (np.linalg.norm(lhs) * np.linalg.norm(rhs)) + 1e-12
    return float(np.dot(lhs, rhs) / denom)


def format_vector_preview(array: np.ndarray, count: int = 8) -> list[float]:
    flat = array.reshape(-1)
    return [float(v) for v in flat[:count]]


def ensure_parent(path: str | Path) -> Path:
    path = Path(path)
    path.parent.mkdir(parents=True, exist_ok=True)
    return path