# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project

"""Intra-unit frame stacking for the omni duplex video track.

Official MiniCPM-o duplex raises the visual refresh rate without touching the
audio cadence: ``get_video_frame_audio_segments(stack_frames=N)`` samples the
``N-1`` sub-frames captured inside each second (skipping the one that would
duplicate the second's base frame) and tiles them into a *single* composite
image. The duplex loop then feeds ``frame_list=[base, composite]`` alongside
that second's audio, so a unit still carries one second of audio and at most
two images however large ``N`` is.

Tiling therefore belongs to media preprocessing, not to the wire: a caller
turns its sub-frames into one composite here and sends it as the second frame
of the append that closes the unit.
"""

from __future__ import annotations

import io
from collections.abc import Sequence
from typing import Any

import pybase64 as base64

_BG_COLOR = (255, 255, 255)
_LINE_COLOR = (0, 0, 0)
_LINE_WIDTH = 6


def _as_image(frame: Any) -> Any:
    """Coerce a PIL image, encoded bytes, or base64/data-URL string to PIL RGB."""
    from PIL import Image

    if isinstance(frame, Image.Image):
        return frame
    if isinstance(frame, (bytes, bytearray)):
        return Image.open(io.BytesIO(bytes(frame))).convert("RGB")
    if isinstance(frame, str):
        payload = frame.split(",")[-1] if ";base64," in frame else frame
        return Image.open(io.BytesIO(base64.b64decode(payload))).convert("RGB")
    raise TypeError(f"unsupported frame type for stacking: {type(frame)}")


def _grid_shape(images: Sequence[Any], *, cell_w: int, cell_h: int) -> tuple[int, int]:
    """Rows and columns for ``images``, matching official ``concat_images``.

    4 frames tile 2x2; 2 and 3 frames pick the row/column arrangement whose
    canvas is closest to square, so a landscape clip stacks vertically and a
    portrait one horizontally; anything else falls back to a single row.
    """
    count = len(images)
    if count == 4:
        return 2, 2
    if count == 1:
        return 1, 1
    if count not in (2, 3):
        return 1, count

    candidates = [(1, count), (count, 1)]

    def squareness(rows: int, cols: int) -> float:
        width = cols * cell_w + (cols - 1) * _LINE_WIDTH
        height = rows * cell_h + (rows - 1) * _LINE_WIDTH
        return abs(width / max(1, height) - 1.0)

    ratios = [squareness(rows, cols) for rows, cols in candidates]
    if count == 2 and ratios[0] == ratios[1]:
        average_aspect = sum(im.width / max(1, im.height) for im in images) / count
        return (1, 2) if average_aspect >= 1.0 else (2, 1)
    return candidates[min(range(len(candidates)), key=lambda index: ratios[index])]


def _letterbox(image: Any, target_w: int, target_h: int) -> Any:
    from PIL import Image

    image = image.convert("RGB")
    width, height = image.size
    scale = min(target_w / width, target_h / height)
    new_size = (max(1, round(width * scale)), max(1, round(height * scale)))
    resized = image.resize(new_size, Image.Resampling.BICUBIC)
    canvas = Image.new("RGB", (target_w, target_h), _BG_COLOR)
    canvas.paste(resized, ((target_w - new_size[0]) // 2, (target_h - new_size[1]) // 2))
    return canvas


def concat_frames(frames: Sequence[Any]) -> Any:
    """Tile a unit's sub-frames into one PIL image, as official stacking does.

    Cells are the largest input width and height, each frame is letterboxed to
    keep its aspect ratio, and a separator band is drawn only in the interior
    seams so the model can tell the sub-frames apart.
    """
    from PIL import Image

    images = [_as_image(frame) for frame in frames]
    if not images:
        raise ValueError("cannot stack an empty frame list")

    cell_w = max(image.width for image in images)
    cell_h = max(image.height for image in images)
    rows, cols = _grid_shape(images, cell_w=cell_w, cell_h=cell_h)

    canvas = Image.new(
        "RGB",
        (cols * cell_w + (cols - 1) * _LINE_WIDTH, rows * cell_h + (rows - 1) * _LINE_WIDTH),
        _LINE_COLOR,
    )
    for index, image in enumerate(images[: rows * cols]):
        row, col = divmod(index, cols)
        canvas.paste(
            _letterbox(image, cell_w, cell_h),
            (col * (cell_w + _LINE_WIDTH), row * (cell_h + _LINE_WIDTH)),
        )
    return canvas


def concat_frames_b64(frames: Sequence[Any], *, quality: int = 95) -> str:
    """``concat_frames`` re-encoded as a base64 JPEG, ready for the wire."""
    buf = io.BytesIO()
    concat_frames(frames).save(buf, format="JPEG", quality=quality)
    return base64.b64encode(buf.getvalue()).decode("ascii")


def unit_subframe_offsets(stack_frames: int, *, unit_s: float = 1.0) -> list[float]:
    """Sub-frame capture offsets inside a unit, in seconds from its start.

    ``stack_frames=5`` yields 0.2/0.4/0.6/0.8 s for a 1 s unit. Offset 0 is
    skipped because official skips it: it is the unit's base frame, already
    sent as ``frame_list[0]``.
    """
    count = max(1, int(stack_frames))
    return [index / count * unit_s for index in range(1, count)]