processor.py

import asyncio
import os
import shutil
import sys
from pathlib import Path
from typing import Dict

from audio_analysis import analyze_audio_file
from models import JobState, JobStatus


def _get_device() -> str:
    """Return the best available device for Demucs inference."""
    device = os.environ.get("DEMUCS_DEVICE")
    if device:
        return device
    try:
        import torch
        if torch.cuda.is_available():
            return "cuda"
        if hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
            return "mps"
    except ImportError:
        pass
    return "cpu"

_jobs: Dict[str, JobStatus] = {}
_queue: asyncio.Queue | None = None

STEM_NAMES = ("vocals", "drums", "bass", "other")


def init_queue() -> None:
    """Call this from inside a running event loop (FastAPI startup)."""
    global _queue
    _queue = asyncio.Queue()


def get_job_status(job_id: str) -> JobStatus | None:
    return _jobs.get(job_id)


def create_job(job_id: str, original_filename: str | None = None) -> JobStatus:
    status = JobStatus(
        job_id=job_id,
        state=JobState.pending,
        progress=0,
        original_filename=original_filename,
    )
    _jobs[job_id] = status
    return status


async def enqueue_job(
    job_id: str,
    input_path: str,
    output_dir: str,
    trim_start_sec: float | None = None,
    trim_end_sec: float | None = None,
    original_filename: str | None = None,
) -> None:
    if _queue is None:
        raise RuntimeError("Queue not initialized. Call init_queue() in startup.")
    create_job(job_id, original_filename=original_filename)
    try:
        _queue.put_nowait((job_id, input_path, output_dir, trim_start_sec, trim_end_sec))
    except asyncio.QueueFull:
        _jobs.pop(job_id, None)
        raise RuntimeError("Processing queue is full. Please try again later.")


async def _run_subprocess_exec(*args: str) -> tuple[int, bytes, bytes]:
    proc = await asyncio.create_subprocess_exec(
        *args,
        stdout=asyncio.subprocess.PIPE,
        stderr=asyncio.subprocess.PIPE,
    )
    stdout, stderr = await proc.communicate()
    return proc.returncode, stdout, stderr


async def _trim_audio(
    input_file: Path,
    job_id: str,
    trim_start_sec: float | None,
    trim_end_sec: float | None,
) -> Path:
    if trim_start_sec is None and trim_end_sec is None:
        return input_file

    trim_start = trim_start_sec or 0.0
    trimmed_file = input_file.parent / f"{job_id}_trimmed.mp3"

    ffmpeg_cmd = [
        "ffmpeg",
        "-y",
        "-hide_banner",
        "-loglevel",
        "error",
        "-ss",
        str(trim_start),
        "-i",
        str(input_file),
    ]

    if trim_end_sec is not None:
        ffmpeg_cmd.extend(["-to", str(trim_end_sec)])

    ffmpeg_cmd.extend(["-c", "copy", str(trimmed_file)])

    code, _, stderr = await _run_subprocess_exec(*ffmpeg_cmd)
    if code != 0:
        fallback_cmd = [
            "ffmpeg",
            "-y",
            "-hide_banner",
            "-loglevel",
            "error",
            "-ss",
            str(trim_start),
            "-i",
            str(input_file),
        ]
        if trim_end_sec is not None:
            fallback_cmd.extend(["-to", str(trim_end_sec)])
        fallback_cmd.extend(["-codec:a", "libmp3lame", "-b:a", "320k", str(trimmed_file)])

        fallback_code, _, fallback_stderr = await _run_subprocess_exec(*fallback_cmd)
        if fallback_code != 0:
            msg = fallback_stderr.decode(errors="replace")[:1200]
            if not msg:
                msg = stderr.decode(errors="replace")[:1200]
            raise RuntimeError(f"Audio trim failed: {msg}")

    return trimmed_file


async def _mix_instrumental(output_dir: Path) -> Path:
    drums = output_dir / "drums.mp3"
    bass = output_dir / "bass.mp3"
    other = output_dir / "other.mp3"
    instrumental = output_dir / "instrumental.mp3"

    cmd = [
        "ffmpeg",
        "-y",
        "-hide_banner",
        "-loglevel",
        "error",
        "-i",
        str(drums),
        "-i",
        str(bass),
        "-i",
        str(other),
        "-filter_complex",
        "amix=inputs=3:duration=longest:normalize=0",
        "-codec:a",
        "libmp3lame",
        "-b:a",
        "320k",
        str(instrumental),
    ]

    code, _, stderr = await _run_subprocess_exec(*cmd)
    if code != 0:
        msg = stderr.decode(errors="replace")[:1200]
        raise RuntimeError(f"Failed to create instrumental stem: {msg}")

    return instrumental


def _move_demucs_stems(demucs_work_dir: Path, output_dir: Path) -> dict[str, str]:
    stems: dict[str, str] = {}
    output_dir.mkdir(parents=True, exist_ok=True)

    for stem_name in STEM_NAMES:
        source_file = next(demucs_work_dir.rglob(f"{stem_name}.mp3"), None)
        if source_file is None:
            continue
        target_file = output_dir / f"{stem_name}.mp3"
        source_file.replace(target_file)
        stems[stem_name] = target_file.name

    return stems


async def _run_demucs(
    job_id: str,
    input_path: str,
    output_dir: str,
    trim_start_sec: float | None = None,
    trim_end_sec: float | None = None,
) -> None:
    _jobs[job_id].state = JobState.processing
    _jobs[job_id].progress = 10
    _jobs[job_id].stems = None
    _jobs[job_id].bpm = None
    _jobs[job_id].bpm_confidence = None
    _jobs[job_id].musical_key = None
    _jobs[job_id].key_scale = None
    _jobs[job_id].key_confidence = None

    input_file = Path(input_path)
    output_directory = Path(output_dir)
    demucs_work_dir = output_directory.parent / f"demucs_{job_id}"
    trimmed_file: Path | None = None

    try:
        demucs_input = input_file
        if trim_start_sec is not None or trim_end_sec is not None:
            _jobs[job_id].progress = 20
            demucs_input = await _trim_audio(
                input_file=input_file,
                job_id=job_id,
                trim_start_sec=trim_start_sec,
                trim_end_sec=trim_end_sec,
            )
            trimmed_file = demucs_input

        proc = await asyncio.create_subprocess_exec(
            sys.executable,
            "-m",
            "demucs",
            "--device",
            _get_device(),
            "--mp3",
            "--mp3-bitrate",
            "320",
            "-o",
            str(demucs_work_dir),
            str(demucs_input),
            stdout=asyncio.subprocess.PIPE,
            stderr=asyncio.subprocess.PIPE,
        )

        _jobs[job_id].progress = 35
        _, stderr = await proc.communicate()
        _jobs[job_id].progress = 80

        if proc.returncode != 0:
            raise RuntimeError(f"Demucs failed: {stderr[:2000].decode(errors='replace')}")

        stems = _move_demucs_stems(demucs_work_dir=demucs_work_dir, output_dir=output_directory)
        if not stems:
            raise RuntimeError("Demucs did not produce any stem MP3 outputs")

        if all(k in stems for k in ("drums", "bass", "other")):
            instrumental_path = await _mix_instrumental(output_directory)
            stems["instrumental"] = instrumental_path.name

        _jobs[job_id].progress = 92
        original_for_analysis = output_directory / "original.mp3"
        try:
            if original_for_analysis.exists():
                analysis = await asyncio.to_thread(analyze_audio_file, original_for_analysis)
                _jobs[job_id].bpm = analysis.get("bpm")
                _jobs[job_id].bpm_confidence = analysis.get("bpm_confidence")
                _jobs[job_id].musical_key = analysis.get("musical_key")
                _jobs[job_id].key_scale = analysis.get("key_scale")
                _jobs[job_id].key_confidence = analysis.get("key_confidence")
                _jobs[job_id].tala = analysis.get("tala")
                _jobs[job_id].time_signature = analysis.get("time_signature")
        except Exception:
            # Analysis is additive metadata only; stem generation should still succeed.
            pass

        _jobs[job_id].stems = stems
        _jobs[job_id].state = JobState.done
        _jobs[job_id].progress = 100

    except Exception as e:
        _jobs[job_id].state = JobState.error
        _jobs[job_id].error = str(e)[:500]
        raise
    finally:
        shutil.rmtree(demucs_work_dir, ignore_errors=True)
        if trimmed_file and trimmed_file.exists():
            trimmed_file.unlink(missing_ok=True)


async def _worker() -> None:
    if _queue is None:
        raise RuntimeError("Queue not initialized before worker started")

    while True:
        job_id, input_path, output_dir, trim_start_sec, trim_end_sec = await _queue.get()
        try:
            await _run_demucs(
                job_id,
                input_path,
                output_dir,
                trim_start_sec=trim_start_sec,
                trim_end_sec=trim_end_sec,
            )
        except Exception:
            pass
        finally:
            _queue.task_done()


async def start_worker() -> None:
    asyncio.create_task(_worker())