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

def get_audio_file_info(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)

    # Get the audio file info
    audio_info = sf.info(audio_file)
    
    bit_depth = {'PCM_16': 16, 'FLOAT': 32}.get(audio_info.subtype, 0)
    
    # Convert duration to minutes and seconds
    minutes, seconds = divmod(audio_info.duration, 60)
    
    # Convert bitrate to kbps
    speed_in_kbps = audio_info.samplerate * bit_depth / 1000
    
    # Create a table with the audio file info
    info_table = f"""
    
    | Information | Value |
    | :---: | :---: |
    | File Name | {os.path.basename(audio_file)} |
    | Duration | {int(minutes)} minutes - {int(seconds)} seconds |
    | Bitrate | {speed_in_kbps} kbp/s |
    | Audio Channels | {audio_info.channels} |
    | Samples per second | {audio_info.samplerate} Hz |
    
    """
    
    # Return the info table
    return info_table

def split_audio_into_clips(audio_file):
    # Read the audio data from the file
    audio_data, sample_rate = torchaudio.load(audio_file, normalize=True)

    # Calculate the duration of each 5-second segment
    segment_duration = 5 * sample_rate

    # Split the audio into segments
    segments = [audio_data[:, i:i+segment_duration] for i in range(0, audio_data.size(1), segment_duration)]

    # Save each segment to a separate file
    output_files = []
    for i, segment in enumerate(segments):
        output_file = f"output_segment_{i+1}.wav"
        torchaudio.save(output_file, segment, sample_rate)
        output_files.append(output_file)

    return output_files

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()

# Create the Gradio interface
main()