import cv2
import zipfile
import os
import tempfile
import gradio as gr


def extract_frames_to_zip(video_path, num_frames=10):
    cap = cv2.VideoCapture(video_path)
    if not cap.isOpened():
        raise RuntimeError(f"Could not open video: {video_path}")

    frame_count = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
    if frame_count <= 0:
        raise RuntimeError("Could not determine frame count.")

    num_frames = min(num_frames, frame_count)

    if num_frames == 1:
        frame_indices = [0]
    else:
        frame_indices = [
            round(i * (frame_count - 1) / (num_frames - 1))
            for i in range(num_frames)
        ]

    # Create a temp directory for this run
    tmp_dir = tempfile.mkdtemp()
    zip_path = os.path.join(tmp_dir, "frames.zip")

    with zipfile.ZipFile(zip_path, "w", compression=zipfile.ZIP_DEFLATED) as zf:
        for idx, frame_idx in enumerate(frame_indices):
            cap.set(cv2.CAP_PROP_POS_FRAMES, frame_idx)
            ret, frame = cap.read()
            if not ret:
                print(f"Warning: Could not read frame at index {frame_idx}")
                continue

            success, buffer = cv2.imencode(".jpg", frame)
            if not success:
                print(f"Warning: Could not encode frame at index {frame_idx}")
                continue

            filename_in_zip = f"frame_{idx:02d}.jpg"
            zf.writestr(filename_in_zip, buffer.tobytes())

    cap.release()
    return zip_path


def gradio_fn(video_file, num_frames):
    """
    video_file: path to uploaded video (from Gradio Video input)
    num_frames: integer from slider
    """
    if video_file is None:
        raise gr.Error("Please upload a video first.")

    # Gradio passes a dict for Video { 'name': ..., 'data': ... } in some versions,
    # but in newer versions it passes a filepath string. Handle both.
    if isinstance(video_file, dict):
        video_path = video_file.get("name") or video_file.get("data")
    else:
        video_path = video_file

    if not video_path or not os.path.exists(video_path):
        raise gr.Error("Uploaded video file not found on the server.")

    zip_path = extract_frames_to_zip(video_path, int(num_frames))
    return zip_path


with gr.Blocks() as demo:
    gr.Markdown("# Video Frame Extractor\nUpload a video and get N evenly spaced frames as a ZIP.")

    with gr.Row():
        video_input = gr.Video(label="Upload video", sources=["upload"])
        num_frames_input = gr.Slider(
            minimum=2,
            maximum=30,
            value=10,
            step=1,
            label="Number of frames to extract",
        )

    zip_output = gr.File(label="Download frames ZIP")

    run_btn = gr.Button("Extract frames")

    run_btn.click(
        fn=gradio_fn,
        inputs=[video_input, num_frames_input],
        outputs=[zip_output],
    )

if __name__ == "__main__":
    demo.launch()