"""vocal isolation: flag off, unavailable, success, timeout, cleanup (Demucs mocked)."""
import os
import sys
import tempfile
import unittest
from pathlib import Path
from unittest import mock

sys.path.insert(0, str(Path(__file__).resolve().parents[1]))

from pydub import AudioSegment
from pydub.generators import Sine

from core import vocal_isolation as vi
from core.registry import TEMP_DIR


def wav(path, ms=1000):
    Sine(300).to_audio_segment(ms).export(str(path), format="wav")
    return path


class FakeProc:
    """Stands in for demucs: writes <out>/htdemucs/<stem>/vocals.wav, then exits."""
    def __init__(self, args, **kw):
        self.never_ends = FakeProc.never_ends
        out = Path(args[args.index("-o") + 1])
        if not self.never_ends and FakeProc.write:
            d = out / "htdemucs" / Path(args[-1]).stem
            d.mkdir(parents=True)
            wav(d / "vocals.wav")
        self.returncode = None if self.never_ends else FakeProc.rc
    never_ends = False
    write = True
    rc = 0
    def poll(self):
        return self.returncode
    def kill(self):
        self.returncode = -9
    def wait(self):
        return self.returncode


def env(**kw):
    base = {"VOCAL_ISOLATION": "1", "VOCAL_ISOLATION_CMD": "/bin/true"}
    base.update(kw)
    return mock.patch.dict(os.environ, base)


class VocalIsolationTests(unittest.TestCase):
    def setUp(self):
        self.src = wav(Path(tempfile.mkdtemp()) / "in.wav")
        FakeProc.never_ends, FakeProc.write, FakeProc.rc = False, True, 0
        self.w = mock.patch.object(vi, "_weights_present", return_value=True)
        self.w.start()
        self.addCleanup(self.w.stop)

    def test_flag_off(self):
        with mock.patch.dict(os.environ, {"VOCAL_ISOLATION": ""}):
            self.assertFalse(vi.available())
            with self.assertRaises(vi.VocalIsolationError) as c:
                vi.isolate_vocals(self.src)
            self.assertEqual(c.exception.status, 503)

    def test_unavailable_without_weights_or_binary(self):
        with env(), mock.patch.object(vi, "_weights_present", return_value=False):
            self.assertFalse(vi.available())
            with self.assertRaises(vi.VocalIsolationError) as c:
                vi.isolate_vocals(self.src)
            self.assertEqual(c.exception.status, 503)
        with env(VOCAL_ISOLATION_CMD=""), mock.patch("shutil.which", return_value=None):
            self.assertFalse(vi.available())

    def test_success_and_cleanup(self):
        with env(), mock.patch("subprocess.Popen", FakeProc):
            self.assertTrue(vi.available())
            out = vi.isolate_vocals(self.src)
        self.assertTrue(out.exists())
        self.assertEqual(AudioSegment.from_file(str(out)).frame_rate, 24000)
        self.assertTrue(self.src.exists())  # original untouched
        self.assertEqual(list(TEMP_DIR.glob("vocal_*")), [])
        self.assertEqual(list(TEMP_DIR.glob("*_vocals.*")), [])

    def test_timeout_kills_and_cleans(self):
        FakeProc.never_ends = True
        with env(VOCAL_ISOLATION_TIMEOUT="0.3"), mock.patch("subprocess.Popen", FakeProc):
            with self.assertRaises(vi.VocalIsolationError) as c:
                vi.isolate_vocals(self.src)
        self.assertEqual(c.exception.status, 504)
        self.assertEqual(list(TEMP_DIR.glob("vocal_*")), [])

    def test_cancel(self):
        import threading
        FakeProc.never_ends = True
        ev = threading.Event(); ev.set()
        with env(), mock.patch("subprocess.Popen", FakeProc):
            with self.assertRaises(vi.VocalIsolationError) as c:
                vi.isolate_vocals(self.src, ev)
        self.assertEqual(c.exception.status, 499)

    def test_too_long_and_failure_and_busy(self):
        with env(VOCAL_ISOLATION_MAX_SECONDS="0.5"):
            with self.assertRaises(vi.VocalIsolationError) as c:
                vi.isolate_vocals(self.src)
            self.assertEqual(c.exception.status, 413)
        FakeProc.write, FakeProc.rc = False, 1
        with env(), mock.patch("subprocess.Popen", FakeProc):
            with self.assertRaises(vi.VocalIsolationError) as c:
                vi.isolate_vocals(self.src)
            self.assertEqual(c.exception.status, 500)
        self.assertEqual(list(TEMP_DIR.glob("vocal_*")), [])
        with env(), vi._busy:
            with self.assertRaises(vi.VocalIsolationError) as c:
                vi.isolate_vocals(self.src)
            self.assertEqual(c.exception.status, 409)


if __name__ == "__main__":
    unittest.main()