import gradio as gr
import numpy as np
import os
import soundfile as sf
from pydub import AudioSegment
from io import BytesIO

def main():
    # Gradio Interface
    with gr.Blocks() as app:
        gr.Markdown(
            """
            # <div align="center"> diablofx Audio Interval Cutter (BETA) </div>
            Want to [support](https://ko-fi.com/diablofx) me? Or [join AI HUB](https://discord.gg/aihub) for more Help!\n
            """
        )
        with gr.Row():
            with gr.Column():
                audio_input = gr.Audio(type='filepath')
                create_info_butt = gr.Button(value='Get Audio File Info', variant='primary')
                split_audio_butt = gr.Button(value='Split Audio into 5s Clips', variant='success')
            with gr.Column():
                output_markdown = gr.Markdown(value="", visible=True)

        create_info_butt.click(fn=get_audio_file_info, inputs=[audio_input], outputs=[output_markdown])
        split_audio_butt.click(fn=split_audio_into_clips, inputs=[audio_input])

        app.queue(max_size=1022).launch(share=True)

def get_audio_file_info(audio_file):
    # ... (unchanged)

def split_audio_into_clips(audio_file):
    # Read the audio data from the file
    audio_data, sample_rate = sf.read(audio_file)

    # Convert to mono if it's not mono
    if len(audio_data.shape) > 1:
        audio_data = np.mean(audio_data, axis=1)

    # Convert to Pydub AudioSegment
    audio_segment = AudioSegment(audio_data.tobytes(), frame_rate=sample_rate, sample_width=audio_data.dtype.itemsize, channels=1)

    # Split the audio into 5s clips
    clips = [audio_segment[i:i + 5000] for i in range(0, len(audio_segment), 5000)]

    # Save each clip and create download buttons
    download_buttons = []
    for i, clip in enumerate(clips):
        clip_data = clip.raw_data
        clip_io = BytesIO(clip_data)
        clip_name = f"clip_{i+1}.wav"
        sf.write(clip_name, clip_io.read(), sample_rate)
        download_buttons.append((f"Download Clip {i+1}", clip_name))

    return download_buttons

# Create the Gradio interface
main()