import os
import subprocess
import tempfile
import numpy as np
import soundfile as sf
import gradio as gr
from df.enhance import init_df, enhance, load_audio

# ──────────────────────────────────────────────────────────────
# 1. 初始化（首次会下载模型到 ~/.cache/DeepFilterNet ）
#    init_df() 返回 (model, state, cfg)
model, df_state, _ = init_df()           

# ──────────────────────────────────────────────────────────────
# 2. 处理函数：文件 → 去噪 → loudnorm → 返回路径
def process_audio(path: str) -> str:
    tmp = tempfile.mkdtemp()
    wav48 = os.path.join(tmp, "in_48k.wav")
    den   = os.path.join(tmp, "den.wav")
    out   = os.path.join(tmp, "out.wav")

    # ① 转 48 kHz mono
    subprocess.run(
        ["ffmpeg", "-y", "-i", path, "-ar", "48000", "-ac", "1", wav48],
        stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL, check=True
    )

    # ② 加载波形 → 去噪
    noisy, _ = load_audio(wav48, sr=48000)          # numpy 1-D float32
    clean_t  = enhance(model, df_state, noisy)      # torch.Tensor
    clean_np = clean_t.detach().cpu().numpy()       # → numpy
    sf.write(den, clean_np.reshape(-1, 1), 48000,
             format="WAV", subtype="FLOAT")

    # ③ loudnorm（-16 LUFS / -1 dBTP）
    subprocess.run(
        ["ffmpeg", "-y", "-i", den,
         "-af", "loudnorm=I=-16:TP=-1:LRA=11",
         out],
        stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL, check=True
    )
    return out

# ──────────────────────────────────────────────────────────────
# 3. Gradio 界面
demo = gr.Interface(
    fn=process_audio,
    inputs=gr.Audio(type="filepath", label="Upload audio (≤ 2 min)"),
    outputs=gr.Audio(type="filepath", label="Enhanced audio"),
    title="DeepFilterNet-v3 Audio Enhancer",
    description="Upload a noisy WAV/MP3 and get a cleaned, loudness-normalized track."
)

demo.queue(max_size=20).launch()