"""Video processing for the SDK: keyframe extraction, per-frame signing."""
import io
import os
import tempfile

import cv2
import numpy as np
from PIL import Image


def extract_keyframes(video_path: str, fps: int = 1) -> list[tuple[int, bytes]]:
    """Extract keyframes from a video file.

    Returns list of (frame_index, png_bytes).
    """
    cap = cv2.VideoCapture(video_path)
    if not cap.isOpened():
        raise ValueError(f"Could not open video: {video_path}")

    video_fps = cap.get(cv2.CAP_PROP_FPS) or 30.0
    frame_interval = max(1, int(video_fps / fps))
    keyframes = []
    frame_idx = 0

    while True:
        ret, frame = cap.read()
        if not ret:
            break
        if frame_idx % frame_interval == 0:
            frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
            img = Image.fromarray(frame_rgb)
            buf = io.BytesIO()
            img.save(buf, format="PNG")
            keyframes.append((frame_idx, buf.getvalue()))
        frame_idx += 1

    cap.release()
    return keyframes


def reassemble_video(
    watermarked_frames: dict[int, bytes],
    input_path: str,
    output_path: str,
) -> None:
    """Replace keyframes in a video with watermarked versions."""
    cap = cv2.VideoCapture(input_path)
    if not cap.isOpened():
        raise ValueError(f"Could not open video: {input_path}")

    fps = cap.get(cv2.CAP_PROP_FPS) or 30.0
    width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
    height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))

    fourcc = cv2.VideoWriter_fourcc(*"mp4v")
    writer = cv2.VideoWriter(output_path, fourcc, fps, (width, height))

    wm_arrays = {}
    for idx, png_bytes in watermarked_frames.items():
        img = Image.open(io.BytesIO(png_bytes)).convert("RGB").resize((width, height))
        wm_arrays[idx] = cv2.cvtColor(np.array(img), cv2.COLOR_RGB2BGR)

    frame_idx = 0
    while True:
        ret, frame = cap.read()
        if not ret:
            break
        if frame_idx in wm_arrays:
            writer.write(wm_arrays[frame_idx])
        else:
            writer.write(frame)
        frame_idx += 1

    cap.release()
    writer.release()
