import os
import sys
import types
import warnings

import numpy as np
import pytest

AUDIO_META_MSG = (
    "`torchaudio.backend.common.AudioMetaData` has been moved to "
    "`torchaudio.AudioMetaData`. Please update the import path."
)


def _stub_df_modules(calls: list[str]) -> tuple[types.ModuleType, list[types.ModuleType]]:
    pkg = types.ModuleType("df")
    mods = []
    for full in ("df.utils", "df.logger", "df.io"):
        mod = types.ModuleType(full)

        def probe(name: str):
            def fn(*_args):
                calls.append(name)
                return "deadbeef"

            return fn

        for fn_name in ("get_git_root", "get_commit_hash", "get_branch_name"):
            setattr(mod, fn_name, probe(f"{full}.{fn_name}"))
        setattr(pkg, full.split(".")[1], mod)
        mods.append(mod)
    return pkg, mods


def test_dfn_shim_neutralizes_git_probes(monkeypatch):
    from producer.engines import denoise_dfn

    calls: list[str] = []
    pkg, mods = _stub_df_modules(calls)
    for name, mod in zip(("df", "df.utils", "df.logger", "df.io"), (pkg, *mods), strict=True):
        monkeypatch.setitem(sys.modules, name, mod)
    assert pkg.logger.get_commit_hash() == "deadbeef"
    calls.clear()
    denoise_dfn._shim_df_git()
    for mod in mods:
        for fn_name in ("get_git_root", "get_commit_hash", "get_branch_name"):
            assert mod.__dict__[fn_name]() is None
    assert calls == []


def test_dfn_shim_silences_torchaudio_warning():
    from producer.engines import denoise_dfn

    code = "import warnings\nwarnings.warn(MESSAGE, UserWarning)"
    denoise_dfn._shim_torchaudio_backend()
    with warnings.catch_warnings(record=True) as caught:
        exec(compile(code, "df/io.py", "exec"), {"__name__": "df.io", "MESSAGE": AUDIO_META_MSG})
    assert caught == []
    with warnings.catch_warnings(record=True) as caught:
        exec(
            compile(code, "other/mod.py", "exec"),
            {"__name__": "other.mod", "MESSAGE": AUDIO_META_MSG},
        )
    assert len(caught) == 1


def _metrics_floor(x, sr):
    from producer import meters

    return meters.noise_floor_db(x, sr)


@pytest.mark.slow
def test_dfn3_reduces_noise(sr, noisy_speech):
    pytest.importorskip("torch")
    pytest.importorskip("df")
    from producer.engines import denoise_dfn

    x = noisy_speech[: sr * 4]
    y, eng, _dev = denoise_dfn.denoise(x, sr, 1.0, "cpu")
    assert "dfn" in eng
    assert _metrics_floor(y, sr) < _metrics_floor(x, sr) - 5.0
    assert np.corrcoef(x, y.astype(np.float64))[0, 1] > 0.9


@pytest.mark.slow
def test_zipenhancer_reduces_noise(sr, noisy_speech):
    pytest.importorskip("torch")
    pytest.importorskip("zipenhancer")
    from producer.engines import denoise_zip

    x = noisy_speech[: sr * 4]
    y, eng, _ = denoise_zip.denoise(x, sr, 1.0, "cpu")
    assert "zipenhancer" in eng
    assert _metrics_floor(y, sr) < _metrics_floor(x, sr) - 5.0


@pytest.mark.slow
def test_mossformer2_enhances(sr, noisy_speech):
    pytest.importorskip("torch")
    pytest.importorskip("clearvoice")
    from producer.engines import enhance_mossformer

    x = noisy_speech[: sr * 4]
    y, eng, _ = enhance_mossformer.enhance(x, sr, 1.0, "cpu")
    assert "mossformer2" in eng
    assert np.all(np.isfinite(y))


@pytest.mark.slow
def test_resemble_enhance(sr, noisy_speech):
    if not os.environ.get("PRODUCER_TEST_RESEMBLE"):
        pytest.skip("set PRODUCER_TEST_RESEMBLE=1 to run the isolated-venv generative engine")
    from producer.engines import enhance_resemble

    x = noisy_speech[: sr * 4]
    y, eng, _ = enhance_resemble.enhance(x, sr, 1.0, "cpu")
    assert "resemble" in eng
    assert np.all(np.isfinite(y))