import gradio as gr
from pydub import AudioSegment, silence
import zipfile
import os

def remove_silence(audio):
    # Remove silence
    non_silent_audio = silence.split_on_silence(audio, min_silence_len=1000, silence_thresh=-40)
    return non_silent_audio

def join_audio_segments(segments, segment_duration):
    joined_segments = []
    current_segment = None
    for segment in segments:
        if current_segment is None:
            current_segment = segment
        elif current_segment.duration_seconds < segment_duration:
            current_segment += segment
        else:
            joined_segments.append(current_segment)
            current_segment = segment

    if current_segment is not None:
        if current_segment.duration_seconds < segment_duration:
            if joined_segments:
                joined_segments[-1] += current_segment
            else:
                joined_segments.append(current_segment)
        else:
            joined_segments.append(current_segment)

    return joined_segments

def process_audio(files, segment_duration):
    zip_file_name = "audio_segments.zip"
    with zipfile.ZipFile(zip_file_name, 'w') as zip_file:
        for i, audio_file in enumerate(files):
            audio = AudioSegment.from_file(audio_file.name)

            # Check if the file is in mp3 format
            is_mp3 = False
            if audio_file.name.lower().endswith('.mp3'):
                is_mp3 = True

            # Remove silence
            non_silent_audio = remove_silence(audio)

            segments = []
            for j, segment in enumerate(non_silent_audio):
                if segment.duration_seconds >= segment_duration:
                    segments.extend(segment[0:segment_duration * 1000] for segment in segment[::segment_duration * 1000])
                else:
                    segments.append(segment)

            # Join segments less than the specified duration with the nearest audio segment
            joined_segments = join_audio_segments(segments, segment_duration)

            for k, segment in enumerate(joined_segments):
                segment_file_name = f"segment_{i + 1}_{k + 1}.wav" if not is_mp3 else f"segment_{i + 1}_{k + 1}.mp3"
                segment.export(segment_file_name, format="wav" if not is_mp3 else "mp3")
                zip_file.write(segment_file_name)
                os.remove(segment_file_name)

    return zip_file_name

def split_audio(files, segment_duration):
    try:
        segment_duration = int(segment_duration)
    except ValueError:
        return "Invalid segment duration. Please enter an integer."
    
    if len(files) == 0:
        return "Please upload audio files."
    
    zip_path = process_audio(files, segment_duration)
    return zip_path

# Gradio Interface
with gr.Blocks(theme=gr.themes.Soft(primary_hue=gr.themes.colors.blue, secondary_hue=gr.themes.colors.blue, spacing_size=gr.themes.sizes.spacing_sm, radius_size=gr.themes.sizes.radius_none)) as app:
    gr.Markdown("## Audio Splitter")
    
    with gr.Row():
        audio_files = gr.File(label="Upload Audio Files", file_count="multiple", file_types=["audio"])
        segment_duration = gr.Number(label="Segment Duration (seconds)", value=5)
    
    split_button = gr.Button("Split Audio Files")
    output = gr.File(label="Download Split Audio Segments")

    split_button.click(split_audio, inputs=[audio_files, segment_duration], outputs=output)

app.launch()