import gradio as gr
import ffmpeg
import os

def extract_audio_from_video(video, start_time, end_time, use_time_cutting):
    if not video:
        raise gr.Error("No video file selected")
    
    input_file = video
    output_file = "media/extracted_audio.mp3"

    os.makedirs(os.path.dirname(output_file), exist_ok=True)

    if os.path.exists(output_file):
        os.remove(output_file)

    if use_time_cutting:
        ffmpeg.input(input_file, ss=start_time, to=end_time).output(
            output_file, format="mp3"
        ).run(overwrite_output=True)
    else:
        ffmpeg.input(input_file).output(output_file, format="mp3").run(
            overwrite_output=True
        )

    return output_file


with gr.Blocks() as demo:
    gr.Markdown("<center><h1> Extract Audio </h1></center><br>")

    with gr.Row():
        video_input = gr.Video(label="Upload Video")
        audio_output = gr.Audio(label="Extracted Audio", type="filepath")

    with gr.Row():
        start_time_input = gr.Number(label="Start Time (seconds)", value=0, precision=0)
        end_time_input = gr.Number(label="End Time (seconds)", value=10, precision=0)
        use_time_cutting_input = gr.Checkbox(
            label="Use Time-Based Cutting", value=False
        )

    extract_button = gr.Button("Extract Audio", variant="primary")

    extract_button.click(
        fn=extract_audio_from_video,
        inputs=[video_input, start_time_input, end_time_input, use_time_cutting_input],
        outputs=audio_output,
    )

demo.launch()