"""Codec-based video preprocessing for LlavaOnevision2 (trust_remote_code).

This module is the codec analogue of ``video_processing_llava_onevision2.py``.
It is invoked when a user calls::

    processor(messages=..., video_backend="codec", max_pixels=...)

and is responsible for:

  - Decoding the video and assembling canvas images via ``cv-preinfer``
    (PyPI: ``codec-video-prep``, requires ``ffmpeg`` on PATH).
  - Running the bundled ``Qwen2VLImageProcessor`` on those canvases with a
    pixel budget that is *aligned* to the canvas dimensions (so the
    smart_resize step never desynchronises ``image_grid_thw`` from the
    codec-emitted ``src_patch_position`` array).
  - Producing the per-patch ``patch_positions`` table that
    ``modeling_llava_onevision2.py`` reads for the 2D-MRoPE block layout.

The result is a ``BatchFeature``-shaped dict containing the same keys that
the frame-sampling video path produces (``pixel_values`` /
``image_grid_thw`` / ``patch_positions``), so downstream
``modeling_llava_onevision2.py`` consumes it without changes.
"""

from __future__ import annotations

import hashlib
import json
import os
import shutil
import subprocess
import tempfile
import warnings
from dataclasses import dataclass, field
from pathlib import Path
from typing import Optional

try:
    import fcntl
except ImportError:
    fcntl = None  # type: ignore

import numpy as np
import torch
from PIL import Image


VISION_START = "<|vision_start|>"
VISION_END = "<|vision_end|>"
IMAGE_PAD = "<|image_pad|>"


# ----------------------------------------------------------------- config

@dataclass
class CodecConfig:
    """All knobs for the codec preprocessing pipeline.

    ``max_pixels`` is shared with the image_processor / video_processor pixel
    budget. The processor sets it from the user's ``max_pixels=`` kwarg, so
    canvas size and HF smart_resize budget stay consistent.
    """

    target_canvas: int = 32
    group_size: int = 32
    images_per_group: int = 4
    patch: int = 14
    max_pixels: int = 150000
    min_group_frames: int = 8
    max_group_frames: int = 64
    spatial_mask_mode: str = "off"
    cache_root: Path = field(default_factory=lambda: Path(
        os.getenv(
            "ONLINE_CODEC_CACHE_DIR",
            os.path.join(
                os.getenv("HF_HOME", os.path.expanduser("~/.cache/huggingface")),
                "online_codec",
            ),
        )
    ))
    timeout_seconds: int = int(os.getenv("ONLINE_CODEC_TIMEOUT", "7200"))

    def validate(self) -> None:
        if self.target_canvas <= 0:
            raise ValueError("CodecConfig.target_canvas must be > 0")
        if self.target_canvas % self.images_per_group != 0:
            raise ValueError(
                "CodecConfig.target_canvas must be divisible by images_per_group"
            )
        if self.group_size % self.images_per_group != 0:
            raise ValueError(
                "CodecConfig.group_size must be divisible by images_per_group"
            )

    def num_sampled_frames(self) -> int:
        return (self.target_canvas // self.images_per_group) * self.group_size


# ---------------------------------------------------------- text/position

def _format_timestamp(seconds: float, decimals: int) -> str:
    return f"<{seconds:.{decimals}f} seconds>"


def convert_positions_to_block_layout(
    positions: torch.Tensor, t: int, h: int, w: int, spatial_merge_size: int = 2,
) -> torch.Tensor:
    """Reorder a (T*H*W, 3) patch position table into 2D-MRoPE block layout."""
    sms = int(spatial_merge_size)
    if sms == 1:
        return positions
    total = int(t) * int(h) * int(w)
    indices = torch.arange(total, device=positions.device).view(t, h, w)
    h_m, w_m = int(h) // sms, int(w) // sms
    indices = (
        indices.view(t, h_m, sms, w_m, sms)
        .permute(0, 1, 3, 2, 4).contiguous().view(total)
    )
    return positions[indices]


def codec_positions_for_processor(
    src_positions: np.ndarray, image_grid_thw: torch.Tensor, device: torch.device,
) -> torch.Tensor:
    positions = torch.from_numpy(src_positions).long().to(device)
    expected_total = int(image_grid_thw.prod(dim=1).sum().item())
    if expected_total != positions.shape[0]:
        raise ValueError(
            "codec patch position length mismatch: "
            f"thw_total={expected_total}, positions={positions.shape[0]}"
        )
    chunks, offset = [], 0
    for row in image_grid_thw:
        t, h, w = int(row[0]), int(row[1]), int(row[2])
        n = t * h * w
        chunks.append(convert_positions_to_block_layout(positions[offset: offset + n], t, h, w))
        offset += n
    return torch.cat(chunks, dim=0)


def _timestamp_runs(
    patch_positions: torch.Tensor, fps: float, decimals: int, spatial_merge_size: int = 2,
) -> list[tuple[str, int]]:
    t_values = patch_positions[:, 0]
    unique_t, counts = torch.unique_consecutive(t_values, return_counts=True)
    merge_factor = int(spatial_merge_size) ** 2
    runs = []
    for t_val, count in zip(unique_t.tolist(), counts.tolist()):
        if int(t_val) < 0:
            continue
        token_count = int(count) // merge_factor
        if token_count <= 0:
            continue
        runs.append((_format_timestamp(float(t_val) / float(fps), decimals), token_count))
    return runs


def rewrite_text_with_codec_positions(
    text: str, patch_positions: torch.Tensor, fps: float, decimals: int,
) -> str:
    """Replace the vision span in a chat-template string with codec-aware tokens."""
    parts = []
    for timestamp, token_count in _timestamp_runs(patch_positions, fps, decimals):
        parts.extend([timestamp, VISION_START, IMAGE_PAD * token_count, VISION_END, "\n"])
    vision_text = "".join(parts)
    first_vs, last_ve = text.find(VISION_START), text.rfind(VISION_END)
    if first_vs == -1 or last_ve == -1:
        return text
    tail_start = last_ve + len(VISION_END)
    if tail_start < len(text) and text[tail_start] == "\n":
        tail_start += 1
    return text[:first_vs] + vision_text + text[tail_start:]


def drop_padding_canvases(
    images: list[Image.Image], src_positions: np.ndarray,
) -> tuple[list[Image.Image], np.ndarray, int]:
    """Drop fully-padding canvases (all-negative timestamps) and their patches."""
    n_canvas = len(images)
    if n_canvas == 0:
        return images, src_positions, 0
    total_patches = src_positions.shape[0]
    if total_patches % n_canvas != 0:
        raise ValueError(
            f"src_positions length {total_patches} not divisible by canvas count {n_canvas}"
        )
    ppc = total_patches // n_canvas
    positions = src_positions.reshape(n_canvas, ppc, 3)
    canvas_t = positions[..., 0]
    keep_mask = (canvas_t >= 0).any(axis=1)
    if bool((keep_mask & ~((canvas_t >= 0).all(axis=1))).any()):
        raise ValueError("encountered half-padding canvas; padding is expected to be canvas-granular")
    dropped = int(n_canvas - int(keep_mask.sum()))
    if dropped == 0:
        return images, src_positions, 0
    kept_images = [img for img, keep in zip(images, keep_mask.tolist()) if keep]
    kept_positions = positions[keep_mask].reshape(-1, 3)
    return kept_images, kept_positions, dropped


# ------------------------------------------------------- cv-preinfer driver

def _get_video_total_frames(video_url: str) -> int:
    import cv2
    cap = cv2.VideoCapture(video_url)
    try:
        total = int(cap.get(cv2.CAP_PROP_FRAME_COUNT) or 0)
    finally:
        cap.release()
    return max(1, total)


def _cache_dir_for(video_url: str, cfg: CodecConfig) -> Path:
    raw = (
        f"{video_url}|tc={cfg.target_canvas}|gs={cfg.group_size}"
        f"|ipg={cfg.images_per_group}|patch={cfg.patch}"
        f"|mp={cfg.max_pixels}|mask={cfg.spatial_mask_mode}"
    )
    key = hashlib.md5(raw.encode()).hexdigest()
    return cfg.cache_root / f"{Path(video_url).stem}_{key}"


def _load_codec_result(out_dir: Path) -> dict:
    with open(out_dir / "meta.json", "r", encoding="utf-8") as f:
        meta = json.load(f)
    canvas_files = meta.get("canvas_files")
    if not canvas_files:
        for ext in ("npy", "jpg", "png"):
            hits = sorted(p.name for p in out_dir.glob(f"canvas_*.{ext}"))
            if hits:
                canvas_files = hits
                break
        canvas_files = canvas_files or []
    images = []
    for name in canvas_files:
        fp = out_dir / name
        if name.endswith(".npy"):
            images.append(Image.fromarray(np.load(fp)))
        else:
            images.append(Image.open(fp).convert("RGB"))
    src_positions = np.load(out_dir / "src_patch_position.npy")
    fps = float(meta.get("fps") or 30.0)
    return {"images": images, "src_positions": src_positions, "fps": fps,
            "out_dir": str(out_dir), "meta": meta}


def _run_cv_preinfer(video_url: str, out_dir: Path, cfg: CodecConfig) -> dict:
    tmp_dir = Path(tempfile.mkdtemp(dir=str(cfg.cache_root), prefix=f".tmp_{out_dir.name[:48]}_"))
    num_sampled = min(cfg.num_sampled_frames(), _get_video_total_frames(video_url))
    cmd = [
        "cv-preinfer", "--video", video_url, "--out_dir", str(tmp_dir),
        "--num_sampled_frames", str(num_sampled),
        "--grouping_mode", "readiness",
        "--group_size", str(cfg.group_size),
        "--images_per_group", str(cfg.images_per_group),
        "--patch", str(cfg.patch),
        "--max_pixels", str(cfg.max_pixels),
        "--readiness_sum_threshold", "0",
        "--min_group_frames", str(cfg.min_group_frames),
        "--max_group_frames", str(cfg.max_group_frames),
        "--avoid_keyframes",
        "--canvas_format", "jpg",
    ]
    try:
        result = subprocess.run(cmd, text=True, capture_output=True, timeout=cfg.timeout_seconds)
        if result.returncode != 0:
            detail = (result.stderr or result.stdout)[-2000:]
            raise RuntimeError(f"online codec failed rc={result.returncode}: {detail}")
        if out_dir.exists():
            shutil.rmtree(out_dir)
        tmp_dir.rename(out_dir)
    except Exception:
        shutil.rmtree(tmp_dir, ignore_errors=True)
        raise
    return _load_codec_result(out_dir)


def process_codec_video(video_url: str, cfg: CodecConfig) -> dict:
    """Public entrypoint: video URL + config -> dict(images, src_positions, fps, ...).

    Result is cached on disk under ``cfg.cache_root``; concurrent workers
    coordinate via a flock-protected sentinel.

    Soft-warning behaviour (B-mode):
      - If the video has fewer frames than needed to fill ``target_canvas``,
        we emit a one-time UserWarning describing the shortfall but proceed
        normally (cv-preinfer will produce fewer canvases than requested).
      - If the video is so short that cv-preinfer cannot form a single
        group (``< min_group_frames``), we emit a clearer warning and let
        cv-preinfer's own error propagate.
    """
    cfg.validate()
    out_dir = _cache_dir_for(video_url, cfg)
    if (out_dir / "meta.json").exists() and (out_dir / "src_patch_position.npy").exists():
        return _load_codec_result(out_dir)

    _maybe_warn_short_video(video_url, cfg)

    cfg.cache_root.mkdir(parents=True, exist_ok=True)
    lock_path = cfg.cache_root / f".{out_dir.name}.lock"
    lock_fd = os.open(str(lock_path), os.O_CREAT | os.O_RDWR, 0o644)
    try:
        if fcntl is not None:
            fcntl.flock(lock_fd, fcntl.LOCK_EX)
        if (out_dir / "meta.json").exists() and (out_dir / "src_patch_position.npy").exists():
            return _load_codec_result(out_dir)
        return _run_cv_preinfer(video_url, out_dir, cfg)
    finally:
        try:
            if fcntl is not None:
                fcntl.flock(lock_fd, fcntl.LOCK_UN)
        finally:
            os.close(lock_fd)


def _maybe_warn_short_video(video_url: str, cfg: CodecConfig) -> None:
    """Soft-warn (B-mode) when a video is too short to fill target_canvas.

    Logic:
      * needed_frames  = num_sampled_frames() = (target_canvas/ipg)*group_size
      * usable_frames  = min(needed_frames, total_frames)
      * expected_canv  = (usable_frames // group_size) * images_per_group
    If ``expected_canv < target_canvas`` we warn. If
    ``total_frames < min_group_frames`` we warn more loudly (cv-preinfer
    will fail downstream and that error is allowed to propagate).
    """
    try:
        total_frames = _get_video_total_frames(video_url)
    except Exception:
        return  # don't fail on probe errors; cv-preinfer will report its own

    needed = cfg.num_sampled_frames()
    usable = min(needed, total_frames)
    expected_canv = (usable // cfg.group_size) * cfg.images_per_group

    if total_frames < cfg.min_group_frames:
        warnings.warn(
            f"[codec] video {video_url!r} has only {total_frames} frames "
            f"(< min_group_frames={cfg.min_group_frames}); cv-preinfer cannot "
            f"form even a single group and will error out. Consider lowering "
            f"min_group_frames or using video_backend='frames' for this clip.",
            UserWarning,
            stacklevel=2,
        )
        return

    if expected_canv < cfg.target_canvas:
        warnings.warn(
            f"[codec] video {video_url!r} has {total_frames} frames; with "
            f"group_size={cfg.group_size}, images_per_group={cfg.images_per_group} "
            f"this yields ~{expected_canv} canvas(es) instead of the requested "
            f"target_canvas={cfg.target_canvas}. Inference will proceed with the "
            f"smaller canvas count.",
            UserWarning,
            stacklevel=2,
        )


# ----------------------------------------------------- processor wiring

def codec_image_processor_outputs(
    image_processor, images: list[Image.Image], max_pixels: int,
) -> dict:
    """Run ``Qwen2VLImageProcessor`` on codec canvases without smart_resize-ing.

    The codec emits canvases already aligned to the patch grid. To keep
    ``image_grid_thw`` consistent with ``src_patch_position``:
      - ``max_pixels`` is clamped up to the largest canvas (never shrinks)
      - ``min_pixels`` is clamped down to the smallest canvas (never upscales)

    Without the ``min_pixels`` clamp, ``Qwen2VLImageProcessor``'s default
    ``min_pixels=200704`` would grow any canvas below that threshold,
    producing extra patches and a chunk/index mismatch downstream.
    """
    canvas_pixels = [im.width * im.height for im in images]
    proc_max = max(int(max_pixels), max(canvas_pixels, default=int(max_pixels)))
    proc_min = min(canvas_pixels) if canvas_pixels else 1
    return image_processor(
        images=images, min_pixels=proc_min, max_pixels=proc_max, return_tensors="pt",
    )


__all__ = [
    "CodecConfig",
    "process_codec_video",
    "drop_padding_canvases",
    "codec_positions_for_processor",
    "rewrite_text_with_codec_positions",
    "codec_image_processor_outputs",
    "VISION_START", "VISION_END", "IMAGE_PAD",
]