#!/usr/bin/env python3
"""音频降噪增强：DeepFilterNet2 人声增强。

上传带噪音频（口播、讲座录音等），输出降噪后的 wav。
模型在首次请求时加载并缓存，之后常驻内存。
"""
import shutil
import subprocess
import threading
import uuid

import gradio as gr

_lock = threading.Lock()
_model = None
_df_state = None


def get_model():
    global _model, _df_state
    if _model is None:
        with _lock:
            if _model is None:
                from df.enhance import init_df

                _model, _df_state, _ = init_df()
    return _model, _df_state


def to_wav(audio_path):
    """把任意格式（m4a/mp3/flac 等）经 ffmpeg 转成 48k 单声道 wav。

    df 自带的 torchaudio 2.0 sox 后端解不了 m4a/AAC，这是报错的根源。
    若容器里没有 ffmpeg，则原样返回，让 load_audio 自己试。
    """
    if shutil.which("ffmpeg") is None:
        return audio_path
    out = f"/tmp/conv_{uuid.uuid4().hex}.wav"
    r = subprocess.run(
        ["ffmpeg", "-y", "-v", "error", "-i", audio_path, "-ar", "48000", "-ac", "1", out],
        capture_output=True,
        text=True,
    )
    if r.returncode != 0:
        raise gr.Error(f"音频转码失败：{r.stderr.strip()[-300:]}")
    return out


def denoise(audio_path, atten_lim, mix):
    if not audio_path:
        return None
    from df.enhance import enhance, load_audio, save_audio

    model, df_state = get_model()
    audio, _ = load_audio(to_wav(audio_path), sr=df_state.sr())
    lim = None if atten_lim == 0 else atten_lim
    enhanced = enhance(model, df_state, audio, atten_lim_db=lim)
    # 混回原声：mix=1 为纯降噪结果，调低可找回被吃掉的细节
    n = min(enhanced.shape[-1], audio.shape[-1])
    out_audio = enhanced[..., :n] * mix + audio[..., :n] * (1.0 - mix)
    out = "/tmp/enhanced.wav"
    save_audio(out, out_audio, df_state.sr())
    return out


with gr.Blocks(title="音频降噪增强") as demo:
    gr.Markdown(
        "# 🎧 音频降噪增强\n"
        "基于 DeepFilterNet2 的人声降噪：上传带底噪、环境声的录音，输出干净版 wav（48kHz）。\n\n"
        "适合口播、讲座这类以人声为主的素材。CPU 推理，长音频请按时长耐心等。"
    )
    inp = gr.Audio(label="原始音频", type="filepath")
    atten = gr.Slider(
        0, 40, value=20, step=1,
        label="最大衰减上限 (dB)",
        info="限制模型对每个频点的抑制强度：数值越小越温和、损失越小；0 = 不限制（最强降噪）。默认 20。",
    )
    mix = gr.Slider(
        0, 1, value=1.0, step=0.05,
        label="降噪结果占比",
        info="1 = 纯降噪结果；调低可按比例混回原声，找回被吃掉的细节（相声里的笑声、掌声也适用）。",
    )
    btn = gr.Button("开始降噪", variant="primary")
    out = gr.Audio(label="降噪后", type="filepath")
    btn.click(denoise, inputs=[inp, atten, mix], outputs=out)


if __name__ == "__main__":
    demo.launch(show_error=True)