import os
import sys
import types
import traceback
import tempfile

import torch
import torchaudio

# ---------------------------------------------------------------------------
# Shim 1: torchaudio.backend.common (removed in torchaudio >= 2.2)
# ---------------------------------------------------------------------------
if not hasattr(torchaudio, "backend") or not hasattr(
    getattr(torchaudio, "backend", None), "common"
):
    _backend_mod = types.ModuleType("torchaudio.backend")
    _common_mod  = types.ModuleType("torchaudio.backend.common")
    _common_mod.AudioMetaData = torchaudio.AudioMetaData
    _backend_mod.common = _common_mod
    sys.modules["torchaudio.backend"]        = _backend_mod
    sys.modules["torchaudio.backend.common"] = _common_mod
    torchaudio.backend = _backend_mod

# ---------------------------------------------------------------------------
# Shim 2: patch torch.Tensor.numpy() so CUDA tensors are silently moved to
# CPU first.  DeepFilterNet 0.5.6 calls .numpy() on tensors that may already
# be on CUDA when the model runs under ZeroGPU.
# ---------------------------------------------------------------------------
_original_numpy = torch.Tensor.numpy


def _safe_numpy(self, *args, **kwargs):
    if self.is_cuda:
        return self.detach().cpu().numpy(*args, **kwargs)
    return _original_numpy(self.detach(), *args, **kwargs)


torch.Tensor.numpy = _safe_numpy

# ---------------------------------------------------------------------------
# Now it is safe to import DeepFilterNet
# ---------------------------------------------------------------------------
from df.enhance import enhance, init_df, load_audio, save_audio  # noqa: E402

import gradio as gr  # noqa: E402
import spaces        # noqa: E402


# ---------------------------------------------------------------------------
# Audio processing
# ---------------------------------------------------------------------------
@spaces.GPU(duration=60)
def clean_audio(input_path: str):
    """Remove background noise and echo from an audio file using DeepFilterNet."""

    if not input_path:
        raise gr.Error("Please upload or record an audio file before processing.")
    if not os.path.exists(input_path):
        raise gr.Error("File not found. Please try uploading the audio again.")
    if os.path.getsize(input_path) == 0:
        raise gr.Error("The uploaded file is empty. Please select a valid audio file.")

    try:
        # Initialise fresh every call — df_state contains ring-buffers that
        # must be on the same device as the model.
        model, df_state, _ = init_df()

        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        model  = model.to(device)

        # load_audio() always returns a CPU tensor
        audio, _ = load_audio(input_path, sr=df_state.sr())

        if audio is None or audio.numel() == 0:
            raise gr.Error("The audio file contains no audio data.")

        audio = audio.to(device)

        with torch.no_grad():
            enhanced = enhance(model, df_state, audio)

        # Explicit CPU move before save_audio() → just in case
        enhanced = enhanced.detach().cpu()

        out_dir     = tempfile.mkdtemp()
        base_name   = os.path.splitext(os.path.basename(input_path))[0]
        output_path = os.path.join(out_dir, f"cleaned_{base_name}.wav")
        save_audio(output_path, enhanced, df_state.sr())

        return output_path

    except gr.Error:
        raise
    except Exception as e:
        traceback.print_exc()
        raise gr.Error(f"Error while processing audio: {str(e)}")


# ---------------------------------------------------------------------------
# UI
# ---------------------------------------------------------------------------
with gr.Blocks(title="Audio Noise & Echo Cleaner (DeepFilterNet)") as demo:
    gr.Markdown(
        """
        # 🎧 Audio Noise, Background & Echo Cleaner

        Upload a file or record audio from your microphone — the app will
        remove background noise, unwanted sounds, and echo using the
        **DeepFilterNet** neural network model.

        ⚡ Processing runs on **ZeroGPU** (free tier): a GPU is dynamically
        allocated only for the duration of each call. No paid API is used —
        everything runs locally inside this Space.
        """
    )

    with gr.Row():
        with gr.Column():
            input_audio = gr.Audio(
                type="filepath",
                label="Input Audio (upload or record)",
            )
            clean_btn = gr.Button("🧹 Clean Audio", variant="primary")

        with gr.Column():
            output_audio = gr.Audio(
                label="Cleaned Audio",
                type="filepath",
            )

    clean_btn.click(
        fn=clean_audio,
        inputs=input_audio,
        outputs=output_audio,
        api_name=False,
    )

    gr.Markdown(
        """
        ---
        **Supported formats:** WAV, MP3, FLAC, OGG and most common audio formats.  
        **Note:** processing is limited to 60 seconds of GPU time per request.
        For very long recordings, split the audio into shorter segments first.
        """
    )

if __name__ == "__main__":
    demo.queue().launch(show_api=False)