import gradio as gr
import numpy as np
import scipy.io.wavfile as wavfile
import scipy.fft as fft
import yt_dlp


def fft_derivative(audio_data: np.ndarray, sample_rate: int) -> np.ndarray:
    freq_domain = fft.fft(audio_data)
    freqs = fft.fftfreq(len(audio_data), 1 / sample_rate)
    freq_derivative = 1j * freqs * freq_domain
    time_domain_derivative = fft.ifft(freq_derivative)
    return np.real(time_domain_derivative)


def convert_file(audio_data: tuple[int, np.ndarray]) -> np.ndarray:
    sample_rate, audio = audio_data
    if len(audio.shape) == 2:
        audio = np.mean(audio, axis=1)
    derivative = fft_derivative(audio, sample_rate)
    derivative = np.int16(derivative / np.max(np.abs(derivative)) * 32767)
    return derivative


def process_audio_file(audio_file: str) -> tuple[int, np.ndarray]:
    sample_rate, audio_data = wavfile.read(audio_file)
    return sample_rate, audio_data


def process_youtube_link(youtube_link: str) -> str:
    song_file = "downloaded_song.wav"
    ydl_opts = {
        'outtmpl': song_file.replace(".wav", ""),
        'postprocessors': [{
            'key': 'FFmpegExtractAudio',
            'preferredcodec': 'wav',
        }]
    }
    with yt_dlp.YoutubeDL(ydl_opts) as ydl:
        ydl.download([youtube_link])
    return song_file


def process_audio(audio: tuple[int, np.ndarray]) -> tuple[int, np.ndarray]:
    return audio[0], convert_file(audio)


with gr.Blocks() as app:
    with gr.Column():
        youtube_link = gr.Textbox(label="YouTube Link", value="https://www.youtube.com/watch?v=YmJIccPWnEk")
        download_button = gr.Button("Download")
        audio_input = gr.Audio(label="Upload Audio File", type="numpy")
        process_button = gr.Button("Process Audio")
        audio_output = gr.Audio(label="Derivative Audio Output")


        def download_audio(youtube_link: str):
            song_file = process_youtube_link(youtube_link)
            sample_rate, audio_data = process_audio_file(song_file)
            return sample_rate, audio_data


        def process_audio_input(audio):
            sample_rate, processed_audio = process_audio(audio)
            return sample_rate, processed_audio


        download_button.click(download_audio, inputs=youtube_link, outputs=[audio_input])
        process_button.click(process_audio_input, inputs=audio_input, outputs=audio_output)

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