#!/usr/bin/env python3
"""
WAV音声をDeepFilterNet3で強調（ノイズ抑制）するスクリプト。

使用例:
    python deepfilternet3_enhance.py input.wav output.wav

    # より強い抑制を試す（ポストフィルタ有効）
    python deepfilternet3_enhance.py input.wav output_pf.wav --post-filter

    # 最大抑制量を12 dBに制限し、音声の変質を抑える
    python deepfilternet3_enhance.py input.wav output_lim12.wav --atten-lim-db 12

    # ステレオ等をモノラル化して処理
    python deepfilternet3_enhance.py input.wav output_mono.wav --mono
"""

from __future__ import annotations

import argparse
import sys
import time
import types
from pathlib import Path

import soundfile as sf
import torch
import torchaudio
from torchaudio.functional import resample as torchaudio_resample


def install_deepfilternet_torchaudio_compat() -> None:
    """DeepFilterNet 0.5.6の旧AudioMetaData importを現行TorchAudio上で通す。"""
    module_name = "torchaudio.backend.common"
    if module_name in sys.modules:
        return

    try:
        __import__(module_name)
        return
    except ModuleNotFoundError:
        # df.ioはimport時の型注釈にだけAudioMetaDataを必要とします。
        # 実際の入出力は、このスクリプトのSoundFile実装を使用します。
        common_module = types.ModuleType(module_name)
        common_module.AudioMetaData = object
        backend_module = types.ModuleType("torchaudio.backend")
        backend_module.common = common_module
        sys.modules["torchaudio.backend"] = backend_module
        sys.modules[module_name] = common_module
        if not hasattr(torchaudio, "backend"):
            torchaudio.backend = backend_module


install_deepfilternet_torchaudio_compat()

from df.enhance import enhance, init_df


MODEL_NAME = "DeepFilterNet3"


def resample_audio(audio: torch.Tensor, original_sr: int, target_sr: int) -> torch.Tensor:
    """高品質なsinc補間で音声をリサンプリングする。"""
    if original_sr == target_sr:
        return audio
    return torchaudio_resample(
        audio,
        original_sr,
        target_sr,
        resampling_method="sinc_interp_kaiser",
        lowpass_filter_width=64,
        rolloff=0.9475937167399596,
        beta=14.769656459379492,
    )


def load_wav(input_path: Path, target_sr: int) -> tuple[torch.Tensor, int]:
    """SoundFileでWAVをfloat32の[C, T]テンソルとして読み込む。"""
    samples, original_sr = sf.read(
        str(input_path),
        dtype="float32",
        always_2d=True,
    )
    audio = torch.from_numpy(samples.T.copy())
    audio = resample_audio(audio, int(original_sr), target_sr)
    return audio.contiguous(), int(original_sr)


def save_wav(output_path: Path, audio: torch.Tensor, sample_rate: int) -> None:
    """SoundFileで[C, T]テンソルをPCM16 WAVとして保存する。"""
    samples = audio.detach().cpu().transpose(0, 1).contiguous().numpy()
    sf.write(
        str(output_path),
        samples,
        sample_rate,
        format="WAV",
        subtype="PCM_16",
    )


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="DeepFilterNet3を用いてWAV音声の背景ノイズを抑制します。"
    )
    parser.add_argument("input_wav", type=Path, help="入力WAVファイル")
    parser.add_argument("output_wav", type=Path, help="出力WAVファイル")
    parser.add_argument(
        "--post-filter",
        action="store_true",
        help="ポストフィルタを有効化し、ノイズ抑制をやや強くします。",
    )
    parser.add_argument(
        "--atten-lim-db",
        type=float,
        default=None,
        metavar="DB",
        help=(
            "最大ノイズ抑制量をdBで制限します。"
            "例: 12を指定すると最大12 dBの抑制となります。"
            "未指定時は制限しません。"
        ),
    )
    parser.add_argument(
        "--mono",
        action="store_true",
        help="多チャンネル入力を平均してモノラル化してから処理します。",
    )
    parser.add_argument(
        "--model-rate-output",
        action="store_true",
        help=(
            "出力をDeepFilterNet3のサンプリング周波数（通常48 kHz）で保存します。"
            "未指定時は入力WAVのサンプリング周波数へ戻します。"
        ),
    )
    return parser.parse_args()


def initialize_deepfilternet3(post_filter: bool):
    """
    DeepFilterNet 0.5.6（戻り値3個）と、mainブランチのAPI（戻り値4個）の
    どちらでも動くよう、先頭2個だけを使用します。
    """
    init_result = init_df(
        model_base_dir=MODEL_NAME,
        post_filter=post_filter,
        log_level="INFO",
        log_file=None,
    )
    model = init_result[0]
    df_state = init_result[1]
    return model, df_state


def enhance_wav(
    input_path: Path,
    output_path: Path,
    *,
    post_filter: bool,
    atten_lim_db: float | None,
    mono: bool,
    model_rate_output: bool,
) -> None:
    if not input_path.is_file():
        raise FileNotFoundError(f"入力ファイルが見つかりません: {input_path}")
    if input_path.suffix.lower() != ".wav":
        raise ValueError(f"入力はWAVファイルを指定してください: {input_path}")
    if output_path.suffix.lower() != ".wav":
        raise ValueError(f"出力拡張子は.wavにしてください: {output_path}")
    if atten_lim_db is not None and atten_lim_db <= 0:
        raise ValueError("--atten-lim-db は0より大きい値を指定してください。")

    output_path.parent.mkdir(parents=True, exist_ok=True)

    print(f"モデルを初期化しています: {MODEL_NAME}")
    model, df_state = initialize_deepfilternet3(post_filter)
    model_sr = int(df_state.sr())

    print(f"入力WAVを読み込んでいます: {input_path}")
    # DeepFilterNet3の処理周波数へ自動リサンプリングされます。
    audio, original_sr = load_wav(input_path, model_sr)

    if mono and audio.shape[0] > 1:
        audio = audio.mean(dim=0, keepdim=True)

    channels = int(audio.shape[0])
    duration_sec = audio.shape[-1] / model_sr
    print(
        f"入力: {channels} ch, 元SR={original_sr} Hz, "
        f"処理SR={model_sr} Hz, 長さ={duration_sec:.1f} 秒"
    )

    start = time.perf_counter()
    enhanced = enhance(
        model,
        df_state,
        audio,
        pad=True,  # STFT/ISTFTおよびモデルの遅延を補償
        atten_lim_db=atten_lim_db,
    )
    elapsed = time.perf_counter() - start
    rtf = elapsed / duration_sec if duration_sec > 0 else float("nan")

    enhanced = enhanced.detach().cpu()

    output_sr = model_sr if model_rate_output else original_sr
    if output_sr != model_sr:
        enhanced = resample_audio(enhanced, model_sr, output_sr)

    # PCM16保存時の整数オーバーフローを防止します。
    peak_before = float(enhanced.abs().max()) if enhanced.numel() else 0.0
    clip_limit = 1.0 - (1.0 / 32768.0)
    enhanced = enhanced.clamp(min=-1.0, max=clip_limit)

    save_wav(output_path, enhanced, output_sr)

    print(f"出力: {output_path}")
    print(f"処理時間: {elapsed:.2f} 秒")
    print(f"Real-time factor: {rtf:.3f}（1未満なら実時間より高速）")
    print(f"処理後ピーク値: {peak_before:.4f}")
    if peak_before > 1.0:
        print(
            "注意: 処理後信号のピークが1.0を超えたため、"
            "PCM16保存時にクリッピングしました。",
            file=sys.stderr,
        )


def main() -> int:
    args = parse_args()
    try:
        enhance_wav(
            args.input_wav,
            args.output_wav,
            post_filter=args.post_filter,
            atten_lim_db=args.atten_lim_db,
            mono=args.mono,
            model_rate_output=args.model_rate_output,
        )
        return 0
    except Exception as exc:
        print(f"エラー: {exc}", file=sys.stderr)
        return 1


if __name__ == "__main__":
    raise SystemExit(main())