import wave
import gradio as gr

def get_audio_duration(audio_file):
    """Get the duration of the audio file in seconds."""
    try:
        with wave.open(audio_file, 'rb') as audio:
            frames = audio.getnframes()
            rate = audio.getframerate()
            duration = frames / float(rate)
        return duration
    except wave.Error:
        raise ValueError("Invalid audio file. Please upload a valid .wav file.")

def get_training_info(audio_file):
    """Determine training parameters based on the duration of the audio file."""
    try:
        if not audio_file:
            raise gr.Error("Please provide an audio file :(")
            return "."
        
        duration = get_audio_duration(audio_file)
        with wave.open(audio_file, 'rb') as audio:
            sample_rate = audio.getframerate()

        training_info = {
            (0, 2): (150, 'TITAN'),
            (2, 3): (200, 'TITAN'),
            (3, 5): (250, 'TITAN'),
            (5, 10): (300, 'normal'),
            (10, 25): (500, 'Normal'),
            (25, 45): (700, 'Normal'),
            (45, 60): (1000, 'Normal')
        }

        for (min_duration, max_duration), (epochs, pretrain) in training_info.items():
            if min_duration <= duration < max_duration:
                return f'You should use the **{pretrain}** pretrain with **{epochs}** epochs at **{sample_rate / 1000} kHz** sample rate. Good luck with your training!'
        
        return 'Duration is not within the specified range.'
    except ValueError as e:
        raise gr.Error(str(e))

with gr.Blocks(theme=gr.themes.Base(primary_hue="sky", secondary_hue="blue"), title="RVC TRAINING HELPER") as demo:
    with gr.Tab("Main Settings"):
        audio_input = gr.Audio(type="filepath", label="Your Audio here")
        start_button = gr.Button("Start!")
        output_text = gr.Textbox(scale=3, label="Your Output here")
        start_button.click(get_training_info, inputs=[audio_input], outputs=[output_text])
    
    with gr.Tab("Credits"):
        gr.Markdown("### This code originally by [TheStinger](https://huggingface.co./TheStinger)<br>Edited by Blane187")

    with gr.Tab('Tutorial'):
        gr.Markdown(
            """
            ### Step by Step Guide

            1. Upload an audio file (or record yourself by tapping the mic button).
            2. Click the `Start!` button.

            And you're done!
            """
        )

    demo.launch(debug=True)