stemdeckapp--stemdeck/app/pipeline/collect.py

from __future__ import annotations

import json
import logging
import shutil
import subprocess
import time
from pathlib import Path

import numpy as np
import soundfile as sf

from app.core.config import (
    DEMUCS_MODEL,
    JOB_TTL_SECONDS,
    STEM_NAMES,
    TIMEOUT_FFMPEG,
    ffmpeg_executable,
)
from app.core.models import Job
from app.core.registry import all_jobs as registry_all
from app.core.registry import persist as registry_persist
from app.core.registry import remove as registry_remove
from app.core.registry import set_proc

logger = logging.getLogger("stemdeck.collect")


def _rmtree(path: Path) -> None:
    try:
        shutil.rmtree(path)
    except FileNotFoundError:
        pass
    except Exception:
        logger.warning("failed to remove %s", path, exc_info=True)


def _run_ffmpeg(job: Job, cmd: list[str]) -> bool:
    """Run an ffmpeg command, registering the subprocess with the job
    registry so POST /api/jobs/{id}/cancel can terminate it. Returns
    True on success, False on failure or external termination.

    Without registering the proc, an in-flight ffmpeg amix would block
    cancellation for up to its 300s timeout -- the cancel flag is set
    but the runner can't see it until subprocess.run returns. With
    set_proc, the cancel API can call proc.terminate() directly and
    communicate() returns within ~1s with a non-zero returncode."""
    proc = subprocess.Popen(cmd, stdout=subprocess.DEVNULL, stderr=subprocess.PIPE)
    set_proc(job.id, proc)
    try:
        try:
            _, stderr = proc.communicate(timeout=TIMEOUT_FFMPEG)
        except subprocess.TimeoutExpired:
            proc.kill()
            proc.communicate()
            logger.warning("ffmpeg timed out for job %s", job.id)
            return False
        if proc.returncode != 0:
            tail = (stderr or b"").decode(errors="replace").splitlines()[-3:]
            logger.warning(
                "ffmpeg exit %s for job %s: %s",
                proc.returncode,
                job.id,
                " | ".join(tail) or "(no stderr)",
            )
            return False
        return True
    finally:
        set_proc(job.id, None)


_TERMINAL = frozenset(("done", "error", "cancelled"))


def collect(job: Job, stems_root: Path, job_dir: Path) -> list[str]:
    """Move Demucs-emitted stems into the job's stems/ dir and clean up
    the demucs intermediate dir. Does NOT delete the source download --
    cleanup_source() is called by the runner after any post-processing
    that needs to re-encode the source (e.g. building original.wav)."""
    target_dir = job_dir / "stems"
    target_dir.mkdir(exist_ok=True)
    found: list[str] = []
    for name in STEM_NAMES:
        src = stems_root / f"{name}.wav"
        if src.exists():
            shutil.move(str(src), target_dir / f"{name}.wav")
            found.append(name)
    _rmtree(job_dir / DEMUCS_MODEL)
    if not found:
        raise RuntimeError("no stems produced by demucs")
    return found


def cleanup_source(job_dir: Path) -> None:
    """Delete the source audio file. Called after collect AND after any
    post-processing that re-encodes the source (make_original_track).
    The source is 100-300 MB, so getting rid of it is the bulk of disk
    reclaim per job; only the stems remain."""
    for f in job_dir.glob("source.*"):
        f.unlink(missing_ok=True)


def make_original_track(job: Job, job_dir: Path, stems_dir: Path) -> Path | None:
    """Build the "Original" backing track at stems/original.wav as the
    sum of the stems the user did NOT select. This way the studio can
    play (original + each selected stem) and reconstruct the full song
    without doubling the selected stems -- which is what would happen
    if "original" were the raw source download (drum hits in original
    + isolated drums.wav = drums at 2x amplitude).

    Skipped when the user kept all 6 stems (no complement to mix) or
    when none of the unselected stem WAVs are on disk."""
    unselected = [s for s in STEM_NAMES if s not in job.selected_stems]
    inputs = [stems_dir / f"{name}.wav" for name in unselected]
    inputs = [p for p in inputs if p.exists()]
    if not inputs:
        return None
    out = stems_dir / "original.wav"
    cmd: list[str] = [
        ffmpeg_executable(),
        "-y",
        "-nostdin",
        "-loglevel",
        "error",
    ]
    for p in inputs:
        cmd += ["-i", str(p)]
    if len(inputs) == 1:
        # Single complement stem -- copy as-is so we still produce a
        # canonical mix.wav-shaped output without invoking amix on a
        # 1-input graph (which is a no-op anyway).
        cmd += ["-c:a", "pcm_s16le", str(out)]
    else:
        filter_inputs = "".join(f"[{i}:a]" for i in range(len(inputs)))
        cmd += [
            "-filter_complex",
            f"{filter_inputs}amix=inputs={len(inputs)}:normalize=0",
            "-c:a",
            "pcm_s16le",
            str(out),
        ]
    return out if _run_ffmpeg(job, cmd) else None


def make_selected_mix(job: Job, stems_dir: Path, found: list[str]) -> Path | None:
    """If the user picked a strict subset of stems at submit time,
    sum those stems with ffmpeg amix into mix.wav. Returns the output
    path on success, or None when there's nothing to mix.

    Returns the existing single stem path (no ffmpeg) if exactly one
    stem was selected -- copying it to mix.wav would be 30 MB of
    duplicate data. The caller uses the returned path's name for the
    download URL, so a single-stem selection points the Download Mix
    button directly at the existing stem file.

    amix normalize=0 keeps stem amplitudes as-is. Demucs separations
    sum back to (close to) the original signal, so a 2-stem subset
    fits comfortably below 0 dBFS without normalization headroom."""
    selected = [s for s in job.selected_stems if s in found]
    if not selected:
        return None
    if len(selected) == 1:
        return stems_dir / f"{selected[0]}.wav"
    inputs = [stems_dir / f"{name}.wav" for name in selected]
    out = stems_dir / "mix.wav"
    cmd: list[str] = [
        ffmpeg_executable(),
        "-y",
        "-nostdin",
        "-loglevel",
        "error",
    ]
    for p in inputs:
        cmd += ["-i", str(p)]
    filter_inputs = "".join(f"[{i}:a]" for i in range(len(inputs)))
    cmd += [
        "-filter_complex",
        f"{filter_inputs}amix=inputs={len(inputs)}:normalize=0",
        "-c:a",
        "pcm_s16le",
        str(out),
    ]
    return out if _run_ffmpeg(job, cmd) else None


_PEAK_POINTS = 1500  # matches OVERVIEW_WAVE_POINTS in player.js


def compute_stem_peaks(stems_dir: Path, stem_names: list[str]) -> None:
    """Compute and cache [min, max] waveform peaks for each stem.
    Failure is non-fatal — missing peaks.json degrades to client-side decode."""
    peaks: dict[str, list[list[float]]] = {}
    for name in stem_names:
        path = stems_dir / f"{name}.wav"
        if not path.is_file():
            continue
        try:
            data, _ = sf.read(path, dtype="float32", always_2d=True)
            ch = data[:, 0]
            n = len(ch)
            if n == 0:
                continue
            chunk = max(1, n // _PEAK_POINTS)
            result: list[list[float]] = []
            for i in range(0, n, chunk):
                block = ch[i : i + chunk]
                result.append([float(np.min(block)), float(np.max(block))])
            peaks[name] = result[:_PEAK_POINTS]
        except Exception:
            logger.warning("could not compute peaks for %s/%s", stems_dir.name, name, exc_info=True)

    if not peaks:
        return

    try:
        tmp = stems_dir / "peaks.json.tmp"
        tmp.write_text(json.dumps(peaks), encoding="utf-8")
        tmp.replace(stems_dir / "peaks.json")
    except Exception:
        logger.warning("could not write peaks.json for %s", stems_dir.name, exc_info=True)


def sweep_old_jobs(jobs_dir: Path) -> None:
    """Delete job directories older than JOB_TTL_SECONDS and remove them from
    the in-memory registry. Called hourly from the background sweep loop
    started at app startup.

    Prefers Job.created_at over directory mtime (which can be touched by
    unrelated filesystem events), and never deletes the directory of an
    active (non-terminal) registered job even if its timestamp looks old.
    Falls back to mtime for orphan directories left over from a previous
    server run, since the registry is in-memory only."""
    cutoff = time.time() - JOB_TTL_SECONDS
    if not jobs_dir.is_dir():
        return
    jobs = registry_all()
    removed = False
    for d in jobs_dir.iterdir():
        if not d.is_dir():
            continue
        job = jobs.get(d.name)
        if job is not None:
            if job.status not in _TERMINAL:
                continue  # never delete an active job's working dir
            if job.created_at >= cutoff:
                continue
        elif d.stat().st_mtime >= cutoff:
            continue
        _rmtree(d)
        registry_remove(d.name)
        removed = True
    if removed:
        registry_persist(jobs_dir)