#!/usr/bin/env python3
"""VoiceGrab separator tests (GPU-upgrade.md §6 automation).

Run with the project venv:
    .venv/bin/python tests/test_separators.py
or with pytest:
    .venv/bin/pytest tests/

The AI (Demucs) end-to-end test needs the model cached or a network
connection (~90 MB, first run only) and takes ~15 s on CPU.
"""
from __future__ import annotations

import math
import os
import subprocess
import sys
import tempfile
import wave

HERE = os.path.dirname(os.path.abspath(__file__))
ROOT = os.path.dirname(HERE)
sys.path.insert(0, ROOT)

import numpy as np  # noqa: E402

import voicegrab as vg  # noqa: E402


# ---------------------------------------------------------------------------
# fixtures / helpers
# ---------------------------------------------------------------------------

def _synthetic_voice(n: int, sr: int = 48000) -> np.ndarray:
    """Speech-like signal (130 Hz fundamental + harmonics, syllable AM)."""
    t = np.arange(n) / sr
    sig = np.zeros_like(t)
    for h in range(1, 31):
        sig += (0.6 / h) * np.sin(2 * np.pi * 130.0 * h * t + h * 0.7)
    sig *= 0.35 + 0.65 * (0.5 + 0.5 * np.sin(2 * np.pi * 4.5 * t))
    sig /= np.abs(sig).max()
    rng = np.random.default_rng(7)
    noise = 0.5 * np.cumsum(rng.normal(0, 1, len(t)))
    noise /= np.abs(noise).max()
    return np.clip(0.7 * sig + 0.5 * noise, -1, 1).astype(np.float32)


def _write_test_mp3(path: str, seconds: int = 5) -> str:
    """Create a seconds-long mp3 with ffmpeg (voice-like tone + noise)."""
    cmd = [
        "ffmpeg", "-hide_banner", "-v", "error", "-y",
        "-f", "lavfi", "-i", f"sine=frequency=220:duration={seconds}",
        "-f", "lavfi", "-i", f"anoisesrc=d={seconds}:c=pink:a=0.3",
        "-filter_complex",
        "[0][1]amix=inputs=2:weights=1 0.5,atrim=0:" + str(seconds) + ",asetpts=PTS-STARTPTS[out]",
        "-map", "[out]", "-c:a", "libmp3lame", path,
    ]
    subprocess.run(cmd, check=True, capture_output=True)
    return path


def _mp3_duration(path: str) -> float:
    out = subprocess.run(
        ["ffprobe", "-v", "error", "-show_entries", "format=duration",
         "-of", "default=nw=1:nk=1", path],
        capture_output=True, text=True, check=True).stdout.strip()
    return float(out)


def _mp3_bitrate(path: str) -> int:
    out = subprocess.run(
        ["ffprobe", "-v", "error", "-show_entries", "format=bit_rate",
         "-of", "default=nw=1:nk=1", path],
        capture_output=True, text=True, check=True).stdout.strip()
    return int(out)


def _ai_available() -> bool:
    return vg._ai_deps_present()


# ---------------------------------------------------------------------------
# unit tests
# ---------------------------------------------------------------------------

def test_modes_constants():
    assert vg.MODE_NONE == "none"
    assert vg.MODE_DSP == "dsp"
    assert vg.MODE_AI == "ai"
    assert vg.AI_MODEL_NAME == "htdemucs_ft"


def test_dsp_available_and_shape():
    sep = vg.DSSeparator(intensity=0.8, stationary=True)
    ok, reason = sep.available()
    assert ok, f"DSP should be available in this venv: {reason}"
    y = _synthetic_voice(48000)
    out, sr = sep.separate(y, 48000, lambda m: None)
    assert sr == 48000, "DSP must not change the sample rate"
    assert out.dtype == np.float32
    assert out.ndim == 1 and len(out) == len(y)
    assert np.isfinite(out).all(), "no NaN/Inf allowed"
    assert float(np.sqrt(np.mean(out**2))) > 0, "must not be all-silence"


def test_ai_available_report():
    sep = vg.DemucsSeparator()
    ok, reason = sep.available()
    if _ai_available():
        assert ok
    else:
        assert not ok
        assert "requirements-ai.txt" in reason


def test_ai_graceful_without_torch():
    """DemucsSeparator.available() must degrade gracefully when torch is absent."""
    code = (
        "import sys, types\n"
        "class Blocker:\n"
        "    def find_spec(self, name, path=None, target=None):\n"
        "        if name == 'torch' or name.startswith('torch.'):\n"
        "            raise ModuleNotFoundError('blocked: ' + name)\n"
        "        return None\n"
        "sys.meta_path.insert(0, Blocker())\n"
        "sys.path.insert(0, %r)\n"
        "import voicegrab as vg\n"
        "ok, reason = vg.DemucsSeparator().available()\n"
        "assert not ok, 'should report unavailable when torch is blocked'\n"
        "assert 'requirements-ai.txt' in reason\n"
        "print('graceful-degradation OK')\n"
    ) % ROOT
    r = subprocess.run([sys.executable, "-c", code],
                       capture_output=True, text=True, timeout=120)
    assert r.returncode == 0, r.stderr
    assert "graceful-degradation OK" in r.stdout


def test_ai_pick_device_order():
    if not _ai_available():
        print("  (skipped: torch/demucs not installed)")
        return
    sep = vg.DemucsSeparator(device="cuda")
    assert sep._pick_device() == "cuda", "explicit device must be honored"
    auto = vg.DemucsSeparator()._pick_device()
    assert auto in ("cpu", "cuda", "mps")


def test_ai_separate_unit():
    if not _ai_available():
        print("  (skipped: torch/demucs not installed)")
        return
    sep = vg.DemucsSeparator()
    y = _synthetic_voice(3 * 48000)   # 3 s @ 48 kHz
    out, sr = sep.separate(y, 48000, lambda m: None)
    assert sr == 44100, "Demucs runs at 44.1 kHz"
    assert out.dtype == np.float32 and out.ndim == 1
    assert np.isfinite(out).all()
    expected = 3 * 44100
    assert abs(len(out) - expected) <= 44100 * 0.02, f"duration drift: {len(out)} vs {expected}"
    assert float(np.sqrt(np.mean(out**2))) > 0, "must not be all-silence"


# ---------------------------------------------------------------------------
# regression: full CLI pipeline (GPU-upgrade.md §6.2)
# ---------------------------------------------------------------------------

def test_cli_dsp_regression():
    with tempfile.TemporaryDirectory(prefix="voicegrab-test-") as wd:
        src = os.path.join(wd, "in.mp3")
        out = os.path.join(wd, "out.mp3")
        _write_test_mp3(src, 5)
        r = subprocess.run(
            [sys.executable, os.path.join(ROOT, "voicegrab.py"),
             "--cli", src, "0", "5", out, "--mode", "dsp"],
            capture_output=True, text=True, timeout=300)
        assert r.returncode == 0, r.stderr
        assert os.path.exists(out)
        dur = _mp3_duration(out)
        assert abs(dur - 5.0) <= 0.15, f"duration {dur} not within 5.0 ± 0.15 s"
        br = _mp3_bitrate(out)
        assert 300_000 <= br <= 350_000, f"expected ~320 kbps, got {br}"


def test_cli_ai_regression():
    if not _ai_available():
        print("  (skipped: torch/demucs not installed)")
        return
    with tempfile.TemporaryDirectory(prefix="voicegrab-test-") as wd:
        src = os.path.join(wd, "in.mp3")
        out = os.path.join(wd, "out.mp3")
        _write_test_mp3(src, 5)
        r = subprocess.run(
            [sys.executable, os.path.join(ROOT, "voicegrab.py"),
             "--cli", src, "0", "5", out, "--mode", "ai"],
            capture_output=True, text=True, timeout=600)
        assert r.returncode == 0, r.stderr
        assert os.path.exists(out)
        dur = _mp3_duration(out)
        assert abs(dur - 5.0) <= 0.15, f"duration {dur} not within 5.0 ± 0.15 s"
        br = _mp3_bitrate(out)
        assert 300_000 <= br <= 350_000, f"expected ~320 kbps, got {br}"


def test_cli_none_mode():
    with tempfile.TemporaryDirectory(prefix="voicegrab-test-") as wd:
        src = os.path.join(wd, "in.mp3")
        out = os.path.join(wd, "out.mp3")
        _write_test_mp3(src, 5)
        r = subprocess.run(
            [sys.executable, os.path.join(ROOT, "voicegrab.py"),
             "--cli", src, "0", "5", out, "--mode", "none"],
            capture_output=True, text=True, timeout=300)
        assert r.returncode == 0, r.stderr
        assert abs(_mp3_duration(out) - 5.0) <= 0.15


def test_cli_rejects_bad_mode():
    with tempfile.TemporaryDirectory(prefix="voicegrab-test-") as wd:
        src = os.path.join(wd, "in.mp3")
        _write_test_mp3(src, 5)
        r = subprocess.run(
            [sys.executable, os.path.join(ROOT, "voicegrab.py"),
             "--cli", src, "0", "5", os.path.join(wd, "o.mp3"), "--mode", "bogus"],
            capture_output=True, text=True, timeout=120)
        assert r.returncode == 2
        assert "none | dsp | ai" in r.stdout


# ---------------------------------------------------------------------------

def main() -> int:
    tests = [v for k, v in sorted(globals().items()) if k.startswith("test_")]
    failed = 0
    for t in tests:
        name = t.__name__
        try:
            t()
            print(f"PASS  {name}")
        except AssertionError as exc:
            failed += 1
            print(f"FAIL  {name}: {exc}")
        except Exception as exc:  # noqa: BLE001
            failed += 1
            print(f"ERROR {name}: {type(exc).__name__}: {exc}")
    total = len(tests)
    print(f"\n{total - failed}/{total} passed")
    return 1 if failed else 0


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