import gradio as gr
import os, tempfile, zipfile, multiprocessing
from moviepy.video.io.ffmpeg_tools import ffmpeg_extract_subclip
from moviepy.editor import VideoFileClip

def extract_clip(args):
    in_path, start, end, out_path = args
    ffmpeg_extract_subclip(in_path, start, end, targetname=out_path)
    return out_path

def split_video(video_data, chunks):
    if not video_data: 
        return [None] * 10, None
    tmp = tempfile.mkdtemp()
    in_path = os.path.join(tmp, "input.mp4")
    with open(in_path, "wb") as f: f.write(video_data)
    clip = VideoFileClip(in_path)
    tot = clip.duration
    seg = tot / chunks
    tasks = [(in_path, i * seg, tot if i == chunks - 1 else (i + 1) * seg, os.path.join(tmp, f"chunk_{i+1}.mp4"))
             for i in range(chunks)]
    with multiprocessing.Pool() as pool:
        parts = pool.map(extract_clip, tasks)
    clip.close()
    zip_path = os.path.join(tmp, "chunks.zip")
    with zipfile.ZipFile(zip_path, "w") as z:
        for p in parts:
            z.write(p, os.path.basename(p))
    return parts, zip_path

def split_interface(video_data, chunks):
    parts, zip_path = split_video(video_data, chunks)
    vids = [gr.update(value=parts[i], visible=True) if i < len(parts)
            else gr.update(visible=False) for i in range(10)]
    return vids + [zip_path]

with gr.Blocks() as demo:
    gr.Markdown("# Equal Chunk Video Splitter\nUpload a video, choose the number of chunks, preview them, and download a ZIP.")
    with gr.Row():
        with gr.Column(scale=1):
            vid_in = gr.File(label="Upload Video", file_types=["video"], type="binary")
            chunks_slider = gr.Slider(1, 10, value=5, step=1, label="Chunks")
            split_btn = gr.Button("Split Video")
        with gr.Column(scale=2):
            zip_out = gr.File(label="Download ZIP")
            previews = [gr.Video(label=f"Chunk {i+1}", visible=False, height=150) for i in range(10)]
    split_btn.click(split_interface, [vid_in, chunks_slider], previews + [zip_out])

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