"""
VIDEO MOSAIC MAKER
==================
Grabs 9 equally spaced frames from a video and returns a 1080p mosaic image,
with the video timestamp of each frame stamped in its upper-left corner.

Requirements:
    pip install opencv-python Pillow
"""

import cv2
from PIL import Image, ImageDraw, ImageFont


def _format_timestamp(seconds):
    """Formats a number of seconds as HH:MM:SS (or MM:SS if under an hour)."""
    total_seconds = int(seconds)
    hours, remainder = divmod(total_seconds, 3600)
    minutes, secs = divmod(remainder, 60)
    if hours:
        return f"{hours:02d}:{minutes:02d}:{secs:02d}"
    return f"{minutes:02d}:{secs:02d}"


def _load_font(size):
    """Tries to load a truetype font, falling back to PIL's default bitmap font."""
    for font_name in ("DejaVuSans-Bold.ttf", "Arial.ttf", "arial.ttf"):
        try:
            return ImageFont.truetype(font_name, size)
        except IOError:
            continue
    return ImageFont.load_default()


def _draw_timestamp(image, text, margin=10):
    """Draws `text` in the upper-left corner of `image` with a dark backing
    box behind it so it stays readable over any frame content."""
    draw = ImageDraw.Draw(image, "RGBA")

    font_size = max(16, image.height // 20)
    font = _load_font(font_size)

    bbox = draw.textbbox((0, 0), text, font=font)
    text_w = bbox[2] - bbox[0]
    text_h = bbox[3] - bbox[1]

    box_x0, box_y0 = margin, margin
    box_x1 = box_x0 + text_w + 2 * margin
    box_y1 = box_y0 + text_h + 2 * margin
    draw.rectangle([box_x0, box_y0, box_x1, box_y1], fill=(0, 0, 0, 160))

    draw.text(
        (box_x0 + margin - bbox[0], box_y0 + margin - bbox[1]),
        text,
        font=font,
        fill=(255, 255, 255, 255),
    )
    return image


def video_mosaic(video_path):
    """
    Takes a video file path and returns a 1920x1080 PIL image
    with 9 frames arranged in a 3x3 grid, each stamped with its
    timestamp in the video.
    """

    # --- Open the video ---
    video = cv2.VideoCapture(video_path)
    total_frames = int(video.get(cv2.CAP_PROP_FRAME_COUNT))
    fps = video.get(cv2.CAP_PROP_FPS) or 30.0  # fallback if fps is unreadable

    # --- Calculate which 9 frames to grab (evenly spaced) ---
    frame_positions = [
        int(total_frames * (i + 1) / 10)  # 10%, 20%, ... 90%
        for i in range(9)
    ]

    # --- Extract each frame, paired with its timestamp ---
    frames = []
    for pos in frame_positions:
        video.set(cv2.CAP_PROP_POS_FRAMES, pos)
        success, frame = video.read()
        if success:
            frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)  # BGR → RGB
            timestamp = _format_timestamp(pos / fps)
            frames.append((Image.fromarray(frame_rgb), timestamp))

    video.release()

    # --- Build the 1920x1080 mosaic ---
    COLS, ROWS = 3, 3
    PADDING    = 10
    CANVAS_W   = 1920
    CANVAS_H   = 1080
    cell_w = (CANVAS_W - PADDING * (COLS + 1)) // COLS
    cell_h = (CANVAS_H - PADDING * (ROWS + 1)) // ROWS

    mosaic = Image.new("RGB", (CANVAS_W, CANVAS_H), color=(0, 0, 0))

    for i, (frame, timestamp) in enumerate(frames):
        row, col = divmod(i, COLS)
        x = PADDING + col * (cell_w + PADDING)
        y = PADDING + row * (cell_h + PADDING)

        resized = frame.resize((cell_w, cell_h), Image.LANCZOS)
        _draw_timestamp(resized, timestamp)

        mosaic.paste(resized, (x, y))

    return mosaic


# --- Example usage ---
if __name__ == "__main__":
    image = video_mosaic("video.mp4")
    image.save("mosaic.jpg")