"""Optional vocal isolation (Demucs CLI). Off unless VOCAL_ISOLATION=1; never downloads weights."""
from __future__ import annotations

import os
import shutil
import subprocess
import tempfile
import threading
import time
import uuid
from pathlib import Path

from core.audio import _duration, _to_wav_24k
from core.registry import TEMP_DIR

_WEIGHTS_GLOB = "955717e8*.th"  # htdemucs
_busy = threading.Lock()  # ponytail: one job at a time; a queue only if several users isolate concurrently


class VocalIsolationError(Exception):
    def __init__(self, message: str, status: int = 500):
        super().__init__(message)
        self.status = status


def _cmd() -> str | None:
    return os.environ.get("VOCAL_ISOLATION_CMD") or shutil.which("demucs")


def _weights_present() -> bool:
    home = Path(os.environ.get("TORCH_HOME") or Path.home() / ".cache" / "torch")
    return any((home / "hub" / "checkpoints").glob(_WEIGHTS_GLOB))


def enabled() -> bool:
    return os.environ.get("VOCAL_ISOLATION", "").strip().lower() in {"1", "true", "yes", "on"}


def available() -> bool:
    return enabled() and _cmd() is not None and _weights_present()


def isolate_vocals(path: Path, cancel: threading.Event | None = None) -> Path:
    """Return a new 24 kHz mono WAV with vocals only; `path` is left untouched."""
    if not enabled():
        raise VocalIsolationError("Vocal isolation is turned off (set VOCAL_ISOLATION=1).", 503)
    cmd = _cmd()
    if cmd is None or not _weights_present():
        raise VocalIsolationError("Vocal isolation is not installed (Demucs or its weights are missing).", 503)
    max_s = float(os.environ.get("VOCAL_ISOLATION_MAX_SECONDS", "900"))
    timeout = float(os.environ.get("VOCAL_ISOLATION_TIMEOUT", "900"))
    if _duration(path) > max_s:
        raise VocalIsolationError(f"Audio is longer than {int(max_s // 60)} minutes; trim it first.", 413)
    if not _busy.acquire(blocking=False):
        raise VocalIsolationError("Another vocal isolation is running; try again shortly.", 409)
    work = Path(tempfile.mkdtemp(prefix="vocal_", dir=TEMP_DIR))
    try:
        proc = subprocess.Popen(
            [cmd, "--two-stems=vocals", "-n", "htdemucs", "-o", str(work), str(path)],
            stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL,
        )
        deadline = time.monotonic() + timeout
        while proc.poll() is None:
            if cancel is not None and cancel.is_set():
                proc.kill(); proc.wait()
                raise VocalIsolationError("Vocal isolation was cancelled.", 499)
            if time.monotonic() > deadline:
                proc.kill(); proc.wait()
                raise VocalIsolationError("Vocal isolation timed out.", 504)
            time.sleep(0.2)
        found = next(work.glob("*/*/vocals.*"), None)
        if proc.returncode != 0 or found is None:
            raise VocalIsolationError("Vocal isolation failed.", 500)
        staged = TEMP_DIR / f"{uuid.uuid4().hex}_vocals{found.suffix}"
        shutil.move(str(found), staged)
        try:
            return _to_wav_24k(staged)
        finally:
            staged.unlink(missing_ok=True)
    finally:
        _busy.release()
        shutil.rmtree(work, ignore_errors=True)