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):
    output_files = []
    for input_file in tqdm(input_files, desc="Converting files"):
        # Load the audio file
        audio = AudioSegment.from_file(input_file)

        # Create the output filename
        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)

        # Ensure the output directory exists
        os.makedirs(session_id, exist_ok=True)

        # Export the audio file to the desired format
        audio.export(output_path, format=output_format)
        output_files.append(output_path)

    return output_files


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, progress=gr.Progress(track_tqdm=True)):
    # Generate a unique session ID using timestamp and UUID
    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)}")
    output_files = convert_audio(files, output_format, session_id)
    zip_filename = create_zip(output_files, session_id)
    print(f"Processing completed!")
    return zip_filename


# List of supported audio formats
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)
        format_choice = gr.Dropdown(choices=audio_formats, label="Output Format", value="mp3")
    submit_button = gr.Button("Convert")
    output_file = gr.File(label="Download Converted File")
    submit_button.click(process_files, inputs=[file_input, format_choice], outputs=output_file)

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