import gradio as gr
import librosa
import soundfile as sf
import torchaudio
import torch
import numpy as np
import zipfile
import tempfile
from pathlib import Path

# =========================================================
# Core Resampling Logic
# =========================================================

def resample_audio(y, sr, target_sr, backend):
    if backend == "librosa":
        y_out = librosa.resample(y, orig_sr=sr, target_sr=target_sr)

    elif backend == "soundfile":
        # soundfile is I/O only, librosa used for resampling
        y_out = librosa.resample(y, orig_sr=sr, target_sr=target_sr)

    elif backend == "torchaudio":
        waveform = torch.tensor(y).unsqueeze(0)
        resampler = torchaudio.transforms.Resample(
            orig_freq=sr,
            new_freq=target_sr
        )
        y_out = resampler(waveform).squeeze(0).numpy()

    return y_out, target_sr


# =========================================================
# Single File Processing
# =========================================================

def single_file_process(audio, target_sr, backend):
    y, sr = audio
    y_out, sr_out = resample_audio(y, sr, target_sr, backend)
    return (sr_out, y_out)


# =========================================================
# Batch ZIP Processing
# =========================================================

def batch_process_zip(zip_file, target_sr, backend):
    output_zip_path = tempfile.NamedTemporaryFile(
        delete=False, suffix=".zip"
    ).name

    with zipfile.ZipFile(zip_file, "r") as zin, \
         zipfile.ZipFile(output_zip_path, "w") as zout:

        for file in zin.namelist():
            if not file.lower().endswith((".wav", ".mp3", ".flac", ".ogg")):
                continue

            # Read file from ZIP
            with zin.open(file) as f:
                with tempfile.NamedTemporaryFile(suffix=".wav") as tmp:
                    tmp.write(f.read())
                    tmp.flush()

                    # Load audio
                    y, sr = librosa.load(tmp.name, sr=None, mono=True)

                    # Resample
                    y_out, sr_out = resample_audio(y, sr, target_sr, backend)

                    # Save output
                    out_name = f"{Path(file).stem}_{backend}_{target_sr}.wav"
                    with tempfile.NamedTemporaryFile(suffix=".wav") as out_tmp:
                        sf.write(out_tmp.name, y_out, sr_out)
                        zout.write(out_tmp.name, out_name)

    return output_zip_path


# =========================================================
# Gradio UI
# =========================================================

with gr.Blocks(title="Audio Resampling Studio") as demo:
    gr.Markdown(
        """
        # 🎚 Audio Resampling Studio
        Resample audio using **librosa**, **soundfile**, or **torchaudio**.

        **Features**
        - 🎧 Single-file processing
        - 📦 Batch ZIP processing
        - 🎯 Sample rates: 16k, 22.05k, 44.1k, 48k
        """
    )

    with gr.Tabs():

        # =================================================
        # Single File Tab
        # =================================================
        with gr.Tab("🎧 Single File"):
            gr.Markdown("### Process a single audio file")

            audio_input = gr.Audio(
                type="numpy",
                label="Upload Audio"
            )

            with gr.Row():
                backend = gr.Radio(
                    ["librosa", "soundfile", "torchaudio"],
                    value="librosa",
                    label="Backend"
                )

                target_sr = gr.Dropdown(
                    [16000, 22050, 44100, 48000],
                    value=16000,
                    label="Target Sample Rate (Hz)"
                )

            process_btn = gr.Button("Resample Audio")
            audio_output = gr.Audio(label="Resampled Output")

            process_btn.click(
                fn=single_file_process,
                inputs=[audio_input, target_sr, backend],
                outputs=audio_output
            )

        # =================================================
        # Batch ZIP Tab
        # =================================================
        with gr.Tab("📦 Batch ZIP"):
            gr.Markdown(
                """
                ### Batch ZIP Processing
                Upload a ZIP file containing audio files.
                You will receive a ZIP of **resampled WAV files**.
                """
            )

            zip_input = gr.File(
                label="Upload ZIP",
                file_types=[".zip"]
            )

            with gr.Row():
                backend_zip = gr.Radio(
                    ["librosa", "soundfile", "torchaudio"],
                    value="librosa",
                    label="Backend"
                )

                target_sr_zip = gr.Dropdown(
                    [16000, 22050, 44100, 48000],
                    value=16000,
                    label="Target Sample Rate (Hz)"
                )

            zip_btn = gr.Button("Process ZIP")
            zip_output = gr.File(label="Download Resampled ZIP")

            zip_btn.click(
                fn=batch_process_zip,
                inputs=[zip_input, target_sr_zip, backend_zip],
                outputs=zip_output
            )

    gr.Markdown(
        """
        ---
        **Output format:** WAV  
        **CPU-safe:** Yes (HF Spaces compatible)  
        **Filename format:**  
        `originalname_backend_samplerate.wav`
        """
    )

demo.launch()