import gradio as gr
import os
import zipfile
import shutil
import moviepy.editor as mp

def process_files(audio_file, video_zip, music_file, csv_file):
    # Unzip the video files
    video_folder = "videos"
    with zipfile.ZipFile(video_zip.name, 'r') as zip_ref:
        zip_ref.extractall(video_folder)

    # Read the CSV file and get the order of videos
    import csv
    video_order = []
    with open(csv_file.name, 'r') as csvfile:
        reader = csv.DictReader(csvfile)
        for row in reader:
            video_order.append(os.path.join(video_folder, f"{row['word']}.mp4"))

    # Concatenate the videos based on the CSV order
    video_clips = [mp.VideoFileClip(video) for video in video_order]
    final_video = mp.concatenate_videoclips(video_clips, method="compose")

    # Add music to the final video if provided
    if music_file is not None:
        music = mp.AudioFileClip(music_file.name)
        final_video = final_video.set_audio(music)

    # Add narration audio to the final video if provided
    if audio_file is not None:
        narration = mp.AudioFileClip(audio_file.name)
        final_video = final_video.set_audio(narration)

    # Save the final video
    output_video_path = "final_video.mp4"
    final_video.write_videofile(output_video_path)

    # Cleanup extracted video files
    shutil.rmtree(video_folder)

    return output_video_path

# Define the Gradio interface
interface = gr.Interface(
    fn=process_files,
    inputs=[
        gr.File(label="Upload Narration Audio File (Optional)"),  # Updated
        gr.File(label="Upload Video Zip File"),
        gr.File(label="Upload Music File"),
        gr.File(label="Upload CSV File")
    ],
    outputs=gr.File(label="Download Final Video"),
    title="Video Concatenation Tool"
)

# Launch the Gradio interface
interface.launch()