import gradio as gr
import os
import ffmpeg
import shutil

# Function to split the audio file using FFmpeg
def split_audio_ffmpeg(input_audio, output_directory, chunk_duration=30):
    """
    Split an audio file into smaller chunks using FFmpeg.
    :param input_audio: Path to the input audio file
    :param output_directory: Directory where output chunks will be saved
    :param chunk_duration: Duration of each chunk in seconds
    """
    os.makedirs(output_directory, exist_ok=True)

    try:
        # Get the total duration of the audio file using ffmpeg.probe()
        probe = ffmpeg.probe(input_audio.name, v='error', select_streams='a', show_entries='stream=duration')
        duration = float(probe['streams'][0]['duration'])
        print(f"Audio duration: {duration} seconds")

        # Split the file into chunks
        chunks = []
        for i in range(0, int(duration), chunk_duration):
            chunk_file = os.path.join(output_directory, f"chunk_{i // chunk_duration + 1}.mp3")
            ffmpeg.input(input_audio.name, ss=i, t=chunk_duration).output(chunk_file).run(quiet=True, overwrite_output=True)
            chunks.append(chunk_file)
            print(f"Generated: {chunk_file}")

        # Create a zip archive of the chunks
        zip_filename = '/tmp/audio_chunks.zip'
        shutil.make_archive(zip_filename.replace('.zip', ''), 'zip', output_directory)
        return zip_filename

    except Exception as e:
        print(f"Error during split: {e}")
        return f"Error during file processing: {e}"

# Gradio interface function
def process_audio(input_audio):
    output_directory = "/tmp/output_chunks"
    zip_file = split_audio_ffmpeg(input_audio, output_directory, chunk_duration=30)

    # Check if zip creation was successful
    if "Error" in zip_file:
        return zip_file

    return gr.File(zip_file)

# Gradio interface setup
interface = gr.Interface(
    fn=process_audio,
    inputs=gr.File(label="Upload Audio", file_types=["audio"]),
    outputs=gr.File(label="Download Chunks as Zip"),
    title="Audio Splitter",
    description="Upload an audio file to split it into smaller chunks of 30 seconds each and download as a zip.",
)

# Launch the Gradio interface
if __name__ == "__main__":
    interface.launch()