# app/services/stems.py — Demucs 4-stem separation subprocess wrapper
#
# Mock mode (CF_STEMS_MOCK=1): copies source file to all 4 stem paths.
# Real mode: runs demucs htdemucs model via asyncio.create_subprocess_exec
#   (list-form, no shell=True — not vulnerable to injection).
#
# Demucs output layout: {output_dir}/htdemucs/{track.stem}/{stem}.wav
from __future__ import annotations

import asyncio
import logging
import os
import shutil
from pathlib import Path

logger = logging.getLogger(__name__)

_MOCK = os.environ.get("CF_STEMS_MOCK", "") == "1"
_STEMS = ("vocals", "drums", "bass", "other")


async def separate(source_audio_path: str, output_dir: str) -> dict[str, str]:
    """
    Split source_audio_path into 4 stems.

    Returns a dict of stem name to absolute file path.
    Raises RuntimeError on demucs failure.
    """
    if _MOCK:
        return await _mock_separate(source_audio_path, output_dir)
    return await _real_separate(source_audio_path, output_dir)


async def _real_separate(source_audio_path: str, output_dir: str) -> dict[str, str]:
    track = Path(source_audio_path)
    logger.info("Running demucs on %s -> %s", track.name, output_dir)
    # create_subprocess_exec takes a list of args — no shell interpolation
    proc = await asyncio.create_subprocess_exec(
        "python", "-m", "demucs",
        "--model", "htdemucs",
        "--out", output_dir,
        str(track),
        stdout=asyncio.subprocess.PIPE,
        stderr=asyncio.subprocess.PIPE,
    )
    _, stderr = await proc.communicate()
    if proc.returncode != 0:
        raise RuntimeError(f"demucs failed (exit {proc.returncode}): {stderr.decode()}")
    return _stem_paths(output_dir, track.stem)


async def _mock_separate(source_audio_path: str, output_dir: str) -> dict[str, str]:
    """Mock: copy source to all 4 stems with a short simulated delay."""
    await asyncio.sleep(1.5)
    track = Path(source_audio_path)
    stem_dir = Path(output_dir) / "htdemucs" / track.stem
    stem_dir.mkdir(parents=True, exist_ok=True)
    paths: dict[str, str] = {}
    for stem in _STEMS:
        dest = str(stem_dir / f"{stem}.wav")
        shutil.copy2(source_audio_path, dest)
        paths[stem] = dest
    logger.info("Mock stems: copied %s -> %s", track.name, stem_dir)
    return paths


def _stem_paths(output_dir: str, track_stem: str) -> dict[str, str]:
    """Build expected output paths from demucs's standard directory layout."""
    stem_dir = Path(output_dir) / "htdemucs" / track_stem
    return {s: str(stem_dir / f"{s}.wav") for s in _STEMS}


def make_stems_dir(data_dir: str, chain_id: str, node_id: str) -> str:
    """Standard stems output directory for a node."""
    return str(Path(data_dir) / "chains" / chain_id / "stems" / node_id)