import os
import uuid
import zipfile
import gradio as gr
from tqdm import tqdm
from datetime import datetime
from pydub import AudioSegment

def convert_audio(input_files, output_format, session_id, merge_files, gap_duration):
    output_files = []
    merged_audio = AudioSegment.silent(duration=0)

    for input_file in tqdm(input_files, desc="Converting files"):
        audio = AudioSegment.from_file(input_file)
        base_name = os.path.splitext(os.path.basename(input_file))[0]
        output_filename = f"{base_name}.{output_format}"
        output_path = os.path.join(session_id, output_filename)
        os.makedirs(session_id, exist_ok=True)
        audio.export(output_path, format=output_format)

        if merge_files:
            merged_audio += audio + AudioSegment.silent(duration=gap_duration)

    if merge_files:
        merged_output_path = os.path.join(session_id, f"merged_output.{output_format}")
        merged_audio.export(merged_output_path, format=output_format)
        return merged_output_path

    return [os.path.join(session_id, f) for f in os.listdir(session_id)]

def create_zip(output_files, session_id):
    zip_filename = f"{session_id}.zip"
    with zipfile.ZipFile(zip_filename, 'w') as zipf:
        for file in tqdm(output_files, desc="Creating ZIP"):
            zipf.write(file, os.path.basename(file))
    return zip_filename

def process_files(files, output_format, merge_files, gap_duration, progress=gr.Progress(track_tqdm=True)):
    session_id = datetime.now().strftime("%Y-%m-%d_%H-%M-%S") + "_" + str(uuid.uuid4())[:8]
    print(f"\nStarting conversion process for session: {session_id}")
    print(f"Total files to convert: {len(files)} to {output_format}")

    output_files = convert_audio(files, output_format, session_id, merge_files, gap_duration)

    print("Processing completed!")
    if merge_files:
        return output_files

    zip_filename = create_zip(output_files, session_id)
    return zip_filename

audio_formats = [
    "wav", "flac", "mp3", "ogg", "aac", "m4a", "aiff", "wma", "opus", "ac3",
    "amr", "dts", "mka", "au", "ra", "voc", "iff", "sd2", "wv", "caf",
    "mpc", "tta", "spx", "gsm"
]

with gr.Blocks() as demo:
    gr.Markdown("## Audio File Converter")
    with gr.Row():
        file_input = gr.Files(file_types=["audio"], height=160)
        with gr.Column():
            format_choice = gr.Dropdown(choices=audio_formats, label="Output Format", value="mp3")
            with gr.Row():
                merge_files_checkbox = gr.Checkbox(label="Merge all files into one")
                gap_slider = gr.Slider(minimum=0, maximum=5000, step=100, value=500, label="Gap between files (ms)")
    submit_button = gr.Button("Convert")
    output_file = gr.File(label="Download Converted File")
    submit_button.click(process_files, inputs=[file_input, format_choice, merge_files_checkbox, gap_slider], outputs=output_file)

if __name__ == "__main__":
    demo.launch()