VoiceStudio/backend/api/routers/batch.py

"""Batch dubbing queue — POST videos with settings, process sequentially.

This is a lightweight batch orchestrator. Each job is a dub project that
runs through the same ingest→transcribe→translate→generate pipeline as
a manual dub, but driven by the queue instead of the UI.

The queue is in-memory (lives for the process lifetime). Jobs persist to
the SQLite `jobs` table for history, but the queue itself restarts empty
on backend restart — intentional, since GPU jobs can't be safely resumed.
"""
import os
import uuid
import time
import asyncio
import logging
from typing import Optional, List

from fastapi import APIRouter, File, UploadFile, HTTPException, Form
from pydantic import BaseModel

from core.config import DATA_DIR
from core import failure
from core.logging_utils import log_safe
from core.file_cleanup import FileCleanupError, unlink_if_present

router = APIRouter()
logger = logging.getLogger("omnivoice.batch")

# ── In-memory queue ─────────────────────────────────────────────────────

_queue: asyncio.Queue = None       # Lazily initialised
_worker_task: asyncio.Task = None  # Background consumer
_jobs: dict = {}                   # job_id → status dict


class BatchJobStatus(BaseModel):
    id: str
    status: str  # "queued" | "running" | "done" | "failed" | "cancelled"
    filename: str
    langs: List[str]
    voice_id: Optional[str] = None
    preserve_bg: bool = True
    created_at: float
    started_at: Optional[float] = None
    finished_at: Optional[float] = None
    error: Optional[str] = None
    progress: Optional[dict] = None


def _ensure_queue():
    """Lazy-init the asyncio queue + worker on first use."""
    global _queue, _worker_task
    if _queue is None:
        _queue = asyncio.Queue()
        _worker_task = asyncio.ensure_future(_worker())


async def _worker():
    """Process jobs one at a time from the queue."""
    while True:
        job_id = await _queue.get()
        job = _jobs.get(job_id)
        if not job or job["status"] == "cancelled":
            _queue.task_done()
            continue

        job["status"] = "running"
        job["started_at"] = time.time()
        logger.info("Batch job %s starting: %s", job_id, job["filename"])

        try:
            await _run_batch_pipeline(job_id, job)
            if job["status"] != "cancelled":
                job["status"] = "done"
                job["finished_at"] = time.time()
                logger.info(
                    "Batch job %s completed in %.1fs",
                    job_id, job["finished_at"] - job["started_at"],
                )
        except asyncio.CancelledError:
            # Task cancellation always means SHUTDOWN: the job-level cancel
            # endpoint only flips job["status"] — nothing ever cancels this
            # task to abort a single job. Swallowing the CancelledError here
            # made the worker unkillable (the while-loop re-entered
            # _queue.get() and event-loop teardown hung forever in
            # _cancel_all_tasks waiting on a task that never finishes). Mark
            # the in-flight job, then let the cancellation propagate.
            job["status"] = "cancelled"
            job["finished_at"] = time.time()
            raise
        except Exception as e:
            job["status"] = "failed"
            # plan-04 (#131): guaranteed non-empty, structured reason.
            job["error"] = failure.build_failure(e, stage="batch", include_diagnostic=False)["reason"]
            job["finished_at"] = time.time()
            logger.error("Batch job %s failed: %s", job_id, e, exc_info=True)
        finally:
            _queue.task_done()


def _set_progress(job, stage, percent=0, **extra):
    """Update a job's progress dict."""
    job["progress"] = {"stage": stage, "percent": percent, **extra}


async def _run_batch_pipeline(job_id: str, job: dict):
    """Full batch dub pipeline: extract → transcribe → translate → generate → mix → export."""
    import subprocess

    loop = asyncio.get_running_loop()
    video_path = job["video_path"]
    langs = job["langs"]
    batch_dir = os.path.join(DATA_DIR, "batch", job_id)
    os.makedirs(batch_dir, exist_ok=True)

    # ── 1. Extract audio ──────────────────────────────────────────────
    _set_progress(job, "extract", 0)
    audio_path = os.path.join(batch_dir, "audio.wav")

    from services.ffmpeg_utils import bed_mix_filter, find_ffmpeg
    ffmpeg = find_ffmpeg()

    def _extract():
        subprocess.run(
            [ffmpeg, "-y", "-i", video_path,
             "-vn", "-acodec", "pcm_s16le", "-ar", "22050", "-ac", "1",
             audio_path],
            stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL,
            timeout=300, check=True,
        )
        # Get duration
        result = subprocess.run(
            [ffmpeg, "-i", audio_path],
            stdout=subprocess.PIPE, stderr=subprocess.PIPE,
            timeout=30,
        )
        import re
        match = re.search(r"Duration: (\d+):(\d+):(\d+)\.(\d+)", result.stderr.decode("utf-8", errors="replace"))
        if match:
            h, m, s, cs = match.groups()
            return int(h) * 3600 + int(m) * 60 + int(s) + int(cs) / 100
        return 0.0

    duration = await loop.run_in_executor(None, _extract)
    job["duration"] = duration
    _set_progress(job, "extract", 100)

    if job["status"] == "cancelled":
        return

    # ── 2. Transcribe ─────────────────────────────────────────────────
    _set_progress(job, "transcribe", 0)

    from services.asr_backend import load_active_asr_backend
    from services.model_manager import _gpu_pool, _cpu_pool, run_on_gpu_pool_guarded
    from services.segmentation import (
        segment_transcript, assign_speakers_heuristic,
    )

    def _transcribe():
        # `load_*`, not `get_*`: the plain selector returns engines whose
        # shallow probe passed but whose deep import chain is broken, failing
        # the whole batch job at `.transcribe()` instead of degrading (#1185).
        backend = load_active_asr_backend()
        result = backend.transcribe(audio_path, word_timestamps=True)
        detected_lang = result.get("language", "en")
        segments = segment_transcript(result, duration=duration)
        segments = assign_speakers_heuristic(segments)
        for i, s in enumerate(segments):
            s["id"] = f"s{i:05x}"
            s.setdefault("text_original", s.get("text", ""))
        try:
            backend.unload()
        except Exception:
            pass
        return segments, detected_lang

    # Bound the batch transcribe (#730) so a wedged whisperx/CTranslate2 call
    # can't hold its GPU-pool worker forever and starve the rest of the backend
    # ("can't reach backend"); run_transcribe_guarded also resets the pool on
    # timeout to restore capacity.
    from services.asr_backend import run_transcribe_guarded
    segments, source_lang = await run_transcribe_guarded(_gpu_pool, _transcribe, what="Batch")
    source_lang = (source_lang or "en").split("_")[0][:2].lower()
    job["segments"] = segments
    job["source_lang"] = source_lang
    _set_progress(job, "transcribe", 100, segments_count=len(segments))

    if job["status"] == "cancelled" or not segments:
        if not segments:
            job["error"] = "Transcription produced no segments"
            job["status"] = "failed"
        return

    # ── Engine resolution (issue #312 class) ────────────────────────────
    # Batch used to hardcode VoiceStudio via get_model() regardless of the
    # engine selected in Model Catalogue → Engines. require_cloning only when a
    # specific voice is pinned (job["voice_id"]) — an unpinned job is fine on
    # any active engine. Resolved ONCE for the whole job (every language
    # below shares the same active engine); an uncaught ValueError here
    # propagates to _worker()'s existing except-Exception handling, which
    # already records a structured job failure via core.failure.build_failure.
    from services.tts_backend import resolve_generation_backend
    backend = await resolve_generation_backend(
        require_cloning=bool(job.get("voice_id")),
        cloning_purpose="this batch job's pinned voice",
    )
    sr = backend.sample_rate

    # ── 3. Translate + Generate per language ───────────────────────────
    total_langs = len(langs)
    outputs = {}

    for lang_idx, target_lang in enumerate(langs):
        if job["status"] == "cancelled":
            return

        # ── 3a. Translate ─────────────────────────────────────────────
        _set_progress(
            job, "translate",
            percent=int((lang_idx / total_langs) * 100),
            current_lang=target_lang,
        )

        translated_segments = list(segments)  # copy
        if target_lang != source_lang:
            try:
                def _translate_batch(segs, src, tgt):
                    """Translate segment texts via Google Translate."""
                    from deep_translator import GoogleTranslator
                    TRANSLATE_CODES = {
                        "en": "en", "es": "es", "fr": "fr", "de": "de",
                        "it": "it", "pt": "pt", "ru": "ru", "ja": "ja",
                        "ko": "ko", "zh": "zh-CN", "ar": "ar", "hi": "hi",
                        "tr": "tr", "pl": "pl", "nl": "nl", "sv": "sv",
                    }
                    src_code = TRANSLATE_CODES.get(src, src) or "auto"
                    tgt_code = TRANSLATE_CODES.get(tgt, tgt)
                    translator = GoogleTranslator(source=src_code, target=tgt_code)
                    out = []
                    for s in segs:
                        s_copy = dict(s)
                        text = s.get("text", "").strip()
                        if text:
                            try:
                                s_copy["text"] = translator.translate(text) or text
                            except Exception as e:
                                logger.warning("Translate seg failed: %s", e)
                        out.append(s_copy)
                    return out

                translated_segments = await loop.run_in_executor(
                    _cpu_pool, _translate_batch,
                    segments, source_lang, target_lang,
                )
            except ImportError:
                logger.warning("deep_translator not installed, skipping translation for %s", target_lang)
            except Exception as e:
                logger.warning("Translation failed for %s: %s, using original", target_lang, e)
                translated_segments = segments

        if job["status"] == "cancelled":
            return

        # ── 3b. Generate TTS ──────────────────────────────────────────
        _set_progress(
            job, "generate",
            percent=int((lang_idx / total_langs) * 100),
            current_lang=target_lang,
            current_segment=0,
            total_segments=len(translated_segments),
        )

        from services.audio_dsp import apply_mastering, normalize_audio
        from services.audio_io import atomic_save_wav
        import torch

        total_samples = int(duration * sr)
        full_audio = torch.zeros(1, total_samples)
        total_segs = len(translated_segments)

        for i, seg in enumerate(translated_segments):
            if job["status"] == "cancelled":
                return

            _set_progress(
                job, "generate",
                percent=int(((lang_idx + (i / total_segs)) / total_langs) * 100),
                current_lang=target_lang,
                current_segment=i + 1,
                total_segments=total_segs,
            )

            seg_start = seg.get("start", 0)
            seg_end = seg.get("end", 0)
            seg_duration = seg_end - seg_start
            seg_text = seg.get("text", "").strip()

            if seg_duration <= 0.05 or not seg_text:
                continue

            def _gen(text=seg_text, lang=target_lang, dur=seg_duration):
                # Normalize once at the segment's text→engine choke point —
                # the same pre-pass as /generate and dub_generate's _gen.
                # `lang` is the job's target language code. Pref-gated,
                # idempotent, never raises.
                from services.text_normalization import normalize_for_tts
                text = normalize_for_tts(text, lang)

                ref_audio = None
                ref_text = None

                # Use voice_id if provided
                if job.get("voice_id"):
                    from core.db import db_conn
                    from core.config import VOICES_DIR as _VD
                    with db_conn() as conn:
                        row = conn.execute(
                            "SELECT * FROM voice_profiles WHERE id=?",
                            (job["voice_id"],),
                        ).fetchone()
                    if row:
                        if row["is_locked"] and row["locked_audio_path"]:
                            ref_audio = os.path.join(_VD, row["locked_audio_path"])
                        elif row["ref_audio_path"]:
                            ref_audio = os.path.join(_VD, row["ref_audio_path"])
                        ref_text = row.get("ref_text")

                try:
                    audio_out = backend.generate(
                        text=text, language=lang,
                        ref_audio=ref_audio, ref_text=ref_text,
                        duration=dur, num_step=16,
                        guidance_scale=2.0, speed=1.0,
                        denoise=True, postprocess_output=True,
                    )
                    if not getattr(backend, "applies_own_mastering", False):
                        audio_out = apply_mastering(audio_out, sample_rate=sr)
                    return normalize_audio(audio_out, target_dBFS=-2.0)
                except Exception as e:
                    logger.warning("TTS failed for seg %d (lang=%s): %s", i, lang, e)
                    # #1190: the silence still stands in for the segment (one
                    # bad line shouldn't bin an otherwise good dub), but it is
                    # no longer INVISIBLE — the job carries a warning the UI /
                    # API consumer can see instead of shipping a
                    # finished-looking track with unexplained silence.
                    job.setdefault("warnings", []).append(
                        f"Segment {i + 1} of the {lang} track failed to "
                        f"synthesize and was left silent: {e}"
                    )
                    return torch.zeros(1, int(dur * sr))

            try:
                # Bounded + pool-reset on hang so a wedged batch segment can't
                # starve the GPU pool and brick the backend (#730 class).
                # Budget is the shared length-scaled one (#1190): a long segment
                # on CPU-class hardware no longer dies on the flat 300s.
                from services.model_manager import generate_timeout_s
                audio_tensor = await run_on_gpu_pool_guarded(
                    _gen, what="Batch generate",
                    timeout=generate_timeout_s(seg_text),
                )

                # Fit to slot
                target_samples_seg = int(seg_duration * sr)
                current_samples = audio_tensor.shape[-1]
                if target_samples_seg > current_samples:
                    audio_tensor = torch.nn.functional.pad(
                        audio_tensor, (0, target_samples_seg - current_samples)
                    )
                elif current_samples > target_samples_seg:
                    audio_tensor = audio_tensor[..., :target_samples_seg]

                # Crossfade
                fade_samples = int(0.015 * sr)
                wl = audio_tensor.shape[-1]
                if wl > fade_samples * 2:
                    ramp_up = torch.linspace(0, 1, fade_samples)
                    ramp_down = torch.linspace(1, 0, fade_samples)
                    audio_tensor[0, :fade_samples] *= ramp_up
                    audio_tensor[0, -fade_samples:] *= ramp_down

                s_idx = int(seg_start * sr)
                e_idx = min(s_idx + wl, total_samples)
                full_audio[:, s_idx:e_idx] += audio_tensor[:, :e_idx - s_idx]

            except TimeoutError as e:
                # #1190/#1202: a GPU timeout (or a saturated pool) used to be
                # swallowed into a silent gap in the dubbed track — the user got
                # a finished-looking video with missing speech and no warning,
                # and on a 1-worker host the abandoned job made every later
                # segment likelier to time out too (the "22-chunk batch dies at
                # chunk 3" cascade). Fail the job loudly instead: _worker()'s
                # except-Exception handler records a structured failure the UI
                # surfaces. Non-timeout per-segment errors keep the old
                # degrade-to-gap behaviour, but are now recorded on the job.
                logger.error("Batch TTS seg %d timed out — failing the job: %s", i, e)
                raise RuntimeError(
                    f"Segment {i + 1} of the {target_lang} track did not "
                    f"render, so the dubbed track would have shipped with a "
                    f"silent gap. {e}"
                ) from e
            except Exception as e:
                logger.warning("Batch TTS seg %d failed: %s", i, e)
                job.setdefault("warnings", []).append(
                    f"Segment {i + 1} of the {target_lang} track failed and was "
                    f"left silent: {e}"
                )

        # ── 3c. Save dubbed audio track ───────────────────────────────
        # Invisible provenance mark on the assembled track (#1169), tensor
        # stage, before the WAV write / aac mux — batch dubs used to ship
        # unmarked while the interactive dub pipeline marked every segment.
        # One whole-track embed (chunked internally, #1045) is equivalent to
        # dub_generate's per-segment marks: the 16-bit message repeats
        # throughout. Runs in the GPU pool like generate's finalize; never
        # raises (degrades to unmarked on failure, same as every producer).
        # Dispatched to the dedicated watermark pool, not the GPU pool (#1190):
        # AudioSeal embedding is CPU work that holds no VRAM, and a whole-track
        # embed is long enough that occupying a GPU worker with it stalled the
        # next language's segments on 1-worker hosts.
        from services.watermark import mark_synthetic
        from services.model_manager import get_watermark_pool
        import functools
        full_audio = await loop.run_in_executor(
            get_watermark_pool(),
            functools.partial(mark_synthetic, full_audio, sr,
                              context="batch.dub_track"),
        )

        # Same assembly pattern as dub_generate.py:390 — `full_audio` is a
        # zero-init tensor that gets +='d from torch.cat-style slices, so
        # it can land non-contiguous + out-of-range. Go through the
        # audited + atomic helper to defend against #48 silent corruption
        # and partial-write truncation simultaneously.
        track_path = os.path.join(batch_dir, f"dubbed_{target_lang}.wav")
        atomic_save_wav(track_path, full_audio, sr)

        # ── 3d. Mix with original video ───────────────────────────────
        _set_progress(
            job, "mix",
            percent=int(((lang_idx + 0.8) / total_langs) * 100),
            current_lang=target_lang,
        )

        output_path = os.path.join(batch_dir, f"output_{target_lang}.mp4")

        def _mix(bg=job.get("preserve_bg", True)):
            if bg:
                # Mix dubbed audio with original background
                subprocess.run(
                    [ffmpeg, "-y",
                     "-i", video_path,
                     "-i", track_path,
                     "-filter_complex",
                     bed_mix_filter("0:a", "1:a", out="out", duration="first"),
                     "-map", "0:v", "-map", "[out]",
                     "-c:v", "copy", "-c:a", "aac", "-b:a", "192k",
                     "-shortest", output_path],
                    stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL,
                    timeout=600, check=True,
                )
            else:
                # Replace audio entirely
                subprocess.run(
                    [ffmpeg, "-y",
                     "-i", video_path,
                     "-i", track_path,
                     "-map", "0:v", "-map", "1:a",
                     "-c:v", "copy", "-c:a", "aac", "-b:a", "192k",
                     "-shortest", output_path],
                    stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL,
                    timeout=600, check=True,
                )

        await loop.run_in_executor(None, _mix)
        outputs[target_lang] = output_path

    job["outputs"] = outputs
    _set_progress(job, "done", 100)


# ── Endpoints ───────────────────────────────────────────────────────────

@router.post("/batch/enqueue")
async def enqueue_batch_job(
    video: UploadFile = File(...),
    langs: str = Form("es"),            # comma-separated lang codes
    voice_id: Optional[str] = Form(None),
    preserve_bg: bool = Form(True),
):
    """Enqueue a video for batch dubbing.

    The video is saved to disk and a job is added to the queue.
    Returns the job ID for status polling.
    """
    _ensure_queue()

    job_id = str(uuid.uuid4())[:12]
    lang_list = [l.strip() for l in langs.split(",") if l.strip()]
    if not lang_list:
        raise HTTPException(400, "At least one target language is required")

    # TTS-only install: no ASR model on disk → typed 409 with a download CTA
    # now, instead of accepting the job and having the transcribe stage
    # silently auto-download multi-GB whisper weights (or fail) in the worker.
    from services.asr_backend import asr_model_missing_detail, asr_model_missing_error
    missing = await asyncio.to_thread(asr_model_missing_error)
    if missing is not None:
        raise HTTPException(409, {**missing, "message": asr_model_missing_detail(missing)})

    # Save the uploaded video
    batch_dir = os.path.join(DATA_DIR, "batch")
    os.makedirs(batch_dir, exist_ok=True)
    ext = os.path.splitext(video.filename or "video.mp4")[1] or ".mp4"
    video_path = os.path.join(batch_dir, f"{job_id}{ext}")

    with open(video_path, "wb") as f:
        content = await video.read()
        f.write(content)

    job = {
        "id": job_id,
        "status": "queued",
        "filename": video.filename or f"{job_id}{ext}",
        "video_path": video_path,
        "langs": lang_list,
        "voice_id": voice_id,
        "preserve_bg": preserve_bg,
        "created_at": time.time(),
        "started_at": None,
        "finished_at": None,
        "error": None,
        "progress": None,
    }
    _jobs[job_id] = job
    await _queue.put(job_id)

    logger.info(
        "Batch job %s enqueued (%d target languages)",
        log_safe(job_id), len(lang_list),
    )
    return {"job_id": job_id, "status": "queued", "queue_position": _queue.qsize()}


@router.get("/batch/jobs")
def list_batch_jobs(status: Optional[str] = None, limit: int = 50):
    """List batch jobs, optionally filtered by status."""
    jobs = list(_jobs.values())
    if status:
        if status == "active":
            jobs = [j for j in jobs if j["status"] in ("queued", "running")]
        else:
            jobs = [j for j in jobs if j["status"] == status]
    jobs.sort(key=lambda j: j["created_at"], reverse=True)
    return jobs[:limit]


@router.get("/batch/jobs/{job_id}")
def get_batch_job(job_id: str):
    """Get the status of a specific batch job."""
    job = _jobs.get(job_id)
    if not job:
        raise HTTPException(404, "Job not found")
    return job


@router.post("/batch/jobs/{job_id}/cancel")
def cancel_batch_job(job_id: str):
    """Cancel a queued or running batch job."""
    job = _jobs.get(job_id)
    if not job:
        raise HTTPException(404, "Job not found")
    if job["status"] in ("done", "failed", "cancelled"):
        return {"already": job["status"]}
    job["status"] = "cancelled"
    job["finished_at"] = time.time()
    return {"cancelled": True}


@router.delete("/batch/jobs/{job_id}")
def delete_batch_job(job_id: str):
    """Delete a batch job record and its video file."""
    job = _jobs.get(job_id)
    if not job:
        raise HTTPException(404, "Job not found")
    if job.get("video_path"):
        try:
            unlink_if_present(job["video_path"])
        except FileCleanupError as exc:
            raise HTTPException(
                status_code=500,
                detail="Could not delete the batch video file. Close any app using it and retry.",
            ) from exc
    _jobs.pop(job_id, None)
    return {"deleted": True}


@router.get("/batch/download/{job_id}/{lang}")
def download_batch_output(job_id: str, lang: str):
    """Download a completed batch job's output video for a given language."""
    from fastapi.responses import FileResponse

    job = _jobs.get(job_id)
    if not job:
        raise HTTPException(404, "Job not found")
    if job["status"] != "done":
        raise HTTPException(400, f"Job is {job['status']}, not done")

    outputs = job.get("outputs", {})
    path = outputs.get(lang)
    if not path or not os.path.exists(path):
        raise HTTPException(404, f"No output for language '{lang}'")

    filename = f"{os.path.splitext(job['filename'])[0]}_{lang}.mp4"
    return FileResponse(
        path,
        media_type="video/mp4",
        filename=filename,
    )