pipeline_runner.py

"""
pipeline_runner.py — Mazinger Dubber pipeline orchestrator for HF ZeroGPU Spaces.

Orchestrates all 10 dubbing stages. GPU-heavy stages are split into two separate
@spaces.GPU-decorated functions to stay within the 120s-per-request hard cap.
"""

from __future__ import annotations

import os
import shutil
import tempfile
import time
from typing import Any, Generator

import spaces  # provided by HF ZeroGPU runtime
import re

# ---------------------------------------------------------------------------
# Constants
# ---------------------------------------------------------------------------

YT_PROXY_SPACE = "HeshamHaroon/yt-proxy"
_YT_RE = re.compile(r'(?:youtube\.com/watch\?v=|youtu\.be/|youtube\.com/shorts/)([a-zA-Z0-9_-]{11})')

HF_TOKEN: str = os.environ.get("HF_TOKEN", "")
LLM_BASE_URL: str = "https://router.huggingface.co/v1"
LLM_MODEL_TEXT: str = "Qwen/Qwen2.5-72B-Instruct"
LLM_MODEL_VISION: str = "Qwen/Qwen2.5-VL-7B-Instruct"
BASE_DIR: str = "/tmp/mazinger_output"

STAGE_NAMES: list[str] = [
    "Download",       # 1
    "Transcribe",     # 2
    "Thumbnails",     # 3
    "Describe",       # 4
    "Review",         # 5  (optional — skipped when asr_review=False)
    "Translate",      # 6
    "Resegment",      # 7
    "Synthesize",     # 8
    "Assemble",       # 9
    "Subtitle",       # 10
]

# ---------------------------------------------------------------------------
# Mazinger imports (only available on HF Spaces where the package is installed)
# ---------------------------------------------------------------------------

from mazinger import ProjectPaths, LLMUsageTracker                              # noqa: E402
from mazinger.llm import build_client                                           # noqa: E402
from mazinger import (                                                          # noqa: E402
    download,
    transcribe,
    thumbnails,
    describe,
    review,
    translate,
    resegment,
    tts,
    assemble,
    subtitle,
)
from mazinger.subtitle import SubtitleStyle, download_google_font               # noqa: E402
from mazinger.srt import parse_file as parse_srt                               # noqa: E402
from mazinger import profiles                                                  # noqa: E402

# ---------------------------------------------------------------------------
# Helper: build LLM client
# ---------------------------------------------------------------------------


def _make_client(model: str = LLM_MODEL_TEXT):
    """Return an OpenAI-compatible client pointed at HF Inference Router."""
    return build_client(api_key=HF_TOKEN, base_url=LLM_BASE_URL)


def _is_youtube_url(url: str) -> bool:
    return bool(_YT_RE.search(url))


def _download_via_proxy(url: str, output_path: str) -> str:
    """Download a YouTube video via the yt-proxy helper Space.

    Uses a background job so we can enforce a timeout (proxy Spaces can be
    asleep and take 1-2 min to wake).
    """
    from gradio_client import Client
    import threading

    print(f"[proxy] Connecting to {YT_PROXY_SPACE}...")
    client = Client(YT_PROXY_SPACE, hf_token=HF_TOKEN)
    print(f"[proxy] Connected. Requesting download of {url} ...")

    # Run predict in a thread so we can enforce a hard timeout
    result_box: list = []
    error_box: list = []

    def _run():
        try:
            r = client.predict(url=url, api_name="/download_video")
            result_box.append(r)
        except Exception as exc:
            error_box.append(exc)

    t = threading.Thread(target=_run, daemon=True)
    t.start()
    t.join(timeout=300)  # 5 min max — proxy may need to wake + download

    if error_box:
        raise RuntimeError(f"YouTube proxy download failed: {error_box[0]}")
    if not result_box:
        raise TimeoutError(
            "YouTube proxy download timed out after 5 minutes. "
            "The proxy Space may be sleeping — try again in a minute."
        )

    result = result_box[0]
    print(f"[proxy] Download complete: {result}")
    os.makedirs(os.path.dirname(output_path) or ".", exist_ok=True)
    shutil.copy2(result, output_path)
    return output_path


# ---------------------------------------------------------------------------
# Helper: call with exponential backoff on 429
# ---------------------------------------------------------------------------


def _call_with_retry(fn, *args, max_retries: int = 3, **kwargs):
    """
    Call *fn* with *args*/*kwargs*, retrying up to *max_retries* times on
    HTTP 429 (rate-limit) errors with exponential backoff (5 s / 15 s / 45 s).
    """
    delays = [5, 15, 45]
    last_exc: Exception | None = None

    for attempt in range(max_retries + 1):
        try:
            return fn(*args, **kwargs)
        except Exception as exc:
            # Detect rate-limit errors by status code attribute or message text
            is_rate_limit = (
                getattr(exc, "status_code", None) == 429
                or "429" in str(exc)
                or "rate limit" in str(exc).lower()
                or "too many requests" in str(exc).lower()
            )
            if is_rate_limit and attempt < max_retries:
                wait = delays[attempt]
                print(
                    f"[retry] 429 rate-limit hit — waiting {wait}s "
                    f"(attempt {attempt + 1}/{max_retries})"
                )
                time.sleep(wait)
                last_exc = exc
                continue
            raise

    # Should never reach here, but satisfy the type checker
    raise last_exc  # type: ignore[misc]


# ---------------------------------------------------------------------------
# Helper: audio duration guard
# ---------------------------------------------------------------------------


def _check_audio_duration(audio_path: str, max_seconds: float = 300.0) -> float:
    """
    Return audio duration in seconds.
    Raises ValueError if the file exceeds *max_seconds*.
    Requires the `soundfile` package (included in mazinger[all-qwen]).
    """
    import soundfile as sf  # lazy import — only needed here

    info = sf.info(audio_path)
    duration: float = info.duration
    if duration > max_seconds:
        raise ValueError(
            f"Audio is {duration:.1f}s — exceeds the {max_seconds:.0f}s limit. "
            "Please trim your video before uploading."
        )
    return duration


# ---------------------------------------------------------------------------
# GPU Stage 1 — Transcription  (120 s allocation)
# ---------------------------------------------------------------------------


@spaces.GPU(duration=120)
def _gpu_transcribe(
    audio_path: str,
    output_path: str,
    method: str = "whisperx",
    model: str | None = None,
    language: str | None = None,
) -> str:
    """
    Run WhisperX (or the requested STT method) on *audio_path* with CUDA.

    Returns the path to the written SRT file (*output_path*).
    First invocation may be slow due to model weight downloads.
    """
    transcribe.transcribe(
        audio_path=audio_path,
        output_path=output_path,
        method=method,
        model=model,
        language=language,
        device="cuda",
    )
    return output_path


# ---------------------------------------------------------------------------
# GPU Stage 2 — TTS Synthesis  (120 s allocation)
# ---------------------------------------------------------------------------


@spaces.GPU(duration=120)
def _gpu_synthesize(
    tts_model_name: str,
    voice_sample: str | None,
    voice_script: str | None,
    voice_theme: str | None,
    clone_profile: str | None,
    srt_entries: list[dict],
    output_dir: str,
    target_language: str,
) -> list[dict]:
    """
    Load the TTS model, resolve a voice prompt via one of the supported modes,
    then synthesize all SRT segments into *output_dir*.

    Voice-prompt priority:
      1. clone_profile  → load a pre-built voice profile
      2. voice_sample + voice_script → create a voice prompt from recorded audio
      3. voice_theme    → load a named built-in theme
      4. fallback       → built-in "narrator-m" voice

    Returns the segment_info list from synthesize_segments (dicts with
    ``idx``, ``start``, ``end``, ``target_dur``, ``wav_path``, ``actual_dur``).
    This is required by :func:`assemble.assemble_timeline`.
    """
    # Load TTS model onto GPU
    tts_model = tts.load_model(tts_model_name, device="cuda")

    # Resolve voice prompt — fetch reference audio + transcript, then create wrapper
    if clone_profile:
        ref_audio, script_path = profiles.fetch_profile(clone_profile)
        # fetch_profile returns file PATHS — read the script content
        with open(script_path) as f:
            ref_text = f.read().strip()
        voice_prompt = tts.create_voice_prompt(tts_model, ref_audio, ref_text)
    elif voice_sample and voice_script:
        voice_prompt = tts.create_voice_prompt(tts_model, voice_sample, voice_script)
    elif voice_theme:
        ref_audio, ref_text = profiles.resolve_theme(voice_theme, target_language, device="cuda")
        voice_prompt = tts.create_voice_prompt(tts_model, ref_audio, ref_text)
    else:
        # Auto-clone: voice_sample is set but voice_script is None
        # For auto-clone, ref_text=None is valid — TTS uses the audio sample only
        if voice_sample:
            voice_prompt = tts.create_voice_prompt(tts_model, voice_sample, None)
        else:
            ref_audio, ref_text = profiles.resolve_theme("narrator-m", target_language, device="cuda")
            voice_prompt = tts.create_voice_prompt(tts_model, ref_audio, ref_text)

    # Synthesize all segments — returns segment_info with wav_path, actual_dur, etc.
    segment_info = tts.synthesize_segments(
        model=tts_model,
        voice_prompt=voice_prompt,
        srt_entries=srt_entries,
        output_dir=output_dir,
        language=target_language,
    )

    return segment_info


# ---------------------------------------------------------------------------
# Main pipeline generator
# ---------------------------------------------------------------------------


def run_pipeline(
    source: str,
    target_language: str,
    voice_mode: str,
    voice_theme: str | None = None,
    voice_profile: str | None = None,
    voice_sample: str | None = None,
    voice_script: str | None = None,
    output_type: str = "video",
    embed_subtitles: bool = True,
    subtitle_font: str = "Cairo",
    subtitle_font_size: int = 28,
    source_language: str = "auto",
    slice_start: str = "",
    slice_end: str = "",
    subtitle_position: str = "bottom",
    subtitle_color: str = "white",
    subtitle_bg_alpha: float = 0.6,
    subtitle_outline_width: int = 1,
    subtitle_bold: bool = False,
    subtitle_line_spacing: int = 8,
    subtitle_source: str = "translated",
    asr_review: bool = False,
    tempo_mode: str = "auto",
    max_tempo: float = 1.5,
    words_per_second: float | None = None,
    duration_budget: float | None = None,
    translate_technical_terms: bool = False,
) -> Generator[tuple[int, str, dict[str, Any]], None, None]:
    """
    Orchestrate all 10 dubbing stages for *source* (URL or local path).

    Yields ``(stage_index, log_message, result_dict)`` tuples at each stage.
    Stage indices are 1-based; a final yield with ``done=True`` is emitted
    after all stages complete.
    """
    from pathlib import Path
    import soundfile as sf
    import hashlib

    # ── Upfront source validation ──────────────────────────────────────
    # ZeroGPU blocks all external DNS except HuggingFace services.
    # Only YouTube URLs (via proxy) and uploaded files work.
    source = source.strip()
    is_local = os.path.exists(source)
    is_url = download.is_url(source)
    is_yt = _is_youtube_url(source)

    if not is_local and not is_url:
        # Might be a URL without scheme — try adding https://
        if "youtube.com" in source or "youtu.be" in source:
            source = "https://" + source
            is_url = True
            is_yt = True
        else:
            yield 1, (
                f"[ERROR] Invalid source: '{source}'\n"
                "Please provide a YouTube URL (e.g. https://youtube.com/watch?v=...) "
                "or upload a video/audio file."
            ), {"error": "invalid source"}
            return

    if is_url and not is_yt:
        yield 1, (
            f"[ERROR] Non-YouTube URLs are not supported on this Space.\n"
            f"URL: {source}\n"
            "ZeroGPU blocks external network access. Only YouTube URLs work "
            "(downloaded via proxy). Please use a YouTube link or upload your file directly."
        ), {"error": "unsupported URL"}
        return

    if is_yt:
        match = _YT_RE.search(source)
        if not match:
            yield 1, (
                f"[ERROR] Could not find a valid YouTube video ID in: {source}\n"
                "Expected format: https://youtube.com/watch?v=VIDEO_ID or https://youtu.be/VIDEO_ID"
            ), {"error": "invalid YouTube URL"}
            return
        slug = match.group(1)
    elif is_local:
        try:
            slug = download.slug_from_path(source)
        except Exception:
            slug = hashlib.md5(source.encode()).hexdigest()[:12]
    else:
        slug = hashlib.md5(source.encode()).hexdigest()[:12]

    proj = ProjectPaths(slug, base_dir=BASE_DIR, target_language=target_language)
    proj.ensure_dirs()
    tracker = LLMUsageTracker()
    result_tmp = tempfile.mkdtemp(prefix="mazinger_result_")

    # -----------------------------------------------------------------------
    # Stage 1 — Download
    # -----------------------------------------------------------------------
    stage = 1
    yield stage, f"[{STAGE_NAMES[stage - 1]}] Downloading: {source}", {}

    try:
        if is_yt:
            yield stage, f"[{STAGE_NAMES[stage - 1]}] Downloading via YouTube proxy...", {}
            _download_via_proxy(source, proj.video)
            download.extract_audio(proj.video, proj.audio)
        elif is_local and download.is_audio_file(source):
            download.ingest_local_audio(source, proj.audio)
        elif is_local:
            download.ingest_local_video(source, proj.video, proj.audio)
        else:
            yield stage, f"[{STAGE_NAMES[stage - 1]}] FAILED: unsupported source", {"error": "unsupported"}
            return
    except Exception as exc:
        yield stage, f"[{STAGE_NAMES[stage - 1]}] FAILED: {exc}", {"error": str(exc)}
        return

    # Slice if requested
    if slice_start or slice_end:
        try:
            download.slice_project(
                proj,
                start=slice_start if slice_start else None,
                end=slice_end if slice_end else None,
            )
            yield stage, f"[{STAGE_NAMES[stage - 1]}] Trimmed to {slice_start or 'start'}–{slice_end or 'end'}.", {}
        except Exception as exc:
            yield stage, f"[{STAGE_NAMES[stage - 1]}] Slice warning: {exc}", {}

    try:
        duration = _check_audio_duration(proj.audio, max_seconds=300.0)
    except ValueError as exc:
        yield stage, f"[{STAGE_NAMES[stage - 1]}] REJECTED: {exc}", {"error": str(exc)}
        return

    yield stage, f"[{STAGE_NAMES[stage - 1]}] Done — {duration:.1f}s of audio.", {}

    # -----------------------------------------------------------------------
    # Stage 2 — Transcribe (GPU)
    # -----------------------------------------------------------------------
    stage = 2
    yield stage, f"[{STAGE_NAMES[stage - 1]}] Transcribing (first run downloads ~3GB model)…", {}

    try:
        _gpu_transcribe(proj.audio, proj.source_srt, method="whisperx")
    except TimeoutError:
        yield stage, f"[{STAGE_NAMES[stage - 1]}] GPU timeout — try a shorter clip.", {"error": "timeout"}
        return
    except Exception as exc:
        yield stage, f"[{STAGE_NAMES[stage - 1]}] FAILED: {exc}", {"error": str(exc)}
        return

    yield stage, f"[{STAGE_NAMES[stage - 1]}] Transcription complete.", {}

    # -----------------------------------------------------------------------
    # Stage 3 — Thumbnails
    # -----------------------------------------------------------------------
    stage = 3
    source_srt_text = Path(proj.source_srt).read_text()
    thumb_paths: list[dict] = []

    if Path(proj.video).exists():
        yield stage, f"[{STAGE_NAMES[stage - 1]}] Extracting keyframes…", {}
        try:
            client = _make_client()
            ts = _call_with_retry(
                thumbnails.select_timestamps,
                source_srt_text, client,
                llm_model=LLM_MODEL_TEXT, usage_tracker=tracker,
            )
            thumb_paths = thumbnails.extract_frames(proj.video, ts, proj.thumbnails_dir)
        except Exception as exc:
            yield stage, f"[{STAGE_NAMES[stage - 1]}] WARNING: {exc}", {}
    else:
        yield stage, f"[{STAGE_NAMES[stage - 1]}] No video — skipping.", {}

    yield stage, f"[{STAGE_NAMES[stage - 1]}] Done — {len(thumb_paths)} keyframe(s).", {}

    # -----------------------------------------------------------------------
    # Stage 4 — Describe (vision LLM)
    # -----------------------------------------------------------------------
    stage = 4
    description: dict = {}

    if thumb_paths:
        yield stage, f"[{STAGE_NAMES[stage - 1]}] Analyzing video content…", {}
        try:
            vision_client = _make_client()
            description = _call_with_retry(
                describe.describe_content,
                source_srt_text, thumb_paths, vision_client,
                llm_model=LLM_MODEL_VISION, usage_tracker=tracker,
            )
        except Exception as exc:
            yield stage, f"[{STAGE_NAMES[stage - 1]}] WARNING: {exc}", {}
    else:
        yield stage, f"[{STAGE_NAMES[stage - 1]}] Skipping (no thumbnails).", {}

    yield stage, f"[{STAGE_NAMES[stage - 1]}] Done.", {}

    # -----------------------------------------------------------------------
    # Stage 5 — Review (optional)
    # -----------------------------------------------------------------------
    stage = 5
    if asr_review:
        yield stage, f"[{STAGE_NAMES[stage - 1]}] Reviewing transcription…", {}
        try:
            review_client = _make_client()
            source_srt_text = _call_with_retry(
                review.review_srt,
                source_srt_text, description, review_client,
                llm_model=LLM_MODEL_TEXT,
                source_language=source_language if source_language != "auto" else "auto",
                usage_tracker=tracker,
            )
            Path(proj.reviewed_srt).write_text(source_srt_text)
            yield stage, f"[{STAGE_NAMES[stage - 1]}] Review complete.", {}
        except Exception as exc:
            yield stage, f"[{STAGE_NAMES[stage - 1]}] WARNING: {exc} — using original.", {}
    else:
        yield stage, f"[{STAGE_NAMES[stage - 1]}] Skipped (not enabled).", {}

    # -----------------------------------------------------------------------
    # Stage 6 — Translate
    # -----------------------------------------------------------------------
    stage = 6
    yield stage, f"[{STAGE_NAMES[stage - 1]}] Translating to {target_language}…", {}

    text_client = _make_client()
    translate_kwargs: dict[str, Any] = {}
    if words_per_second is not None:
        translate_kwargs["words_per_second"] = words_per_second
    if duration_budget is not None:
        translate_kwargs["duration_budget"] = duration_budget

    try:
        translated_srt = _call_with_retry(
            translate.translate_srt,
            source_srt_text, description, thumb_paths, text_client,
            llm_model=LLM_MODEL_TEXT,
            source_language=source_language if source_language != "auto" else "auto",
            target_language=target_language,
            translate_technical_terms=translate_technical_terms,
            usage_tracker=tracker,
            **translate_kwargs,
        )
        Path(proj.translated_raw_srt).write_text(translated_srt)
    except Exception as exc:
        yield stage, f"[{STAGE_NAMES[stage - 1]}] FAILED: {exc}", {"error": str(exc)}
        return

    yield stage, f"[{STAGE_NAMES[stage - 1]}] Translation complete.", {}

    # -----------------------------------------------------------------------
    # Stage 7 — Resegment
    # -----------------------------------------------------------------------
    stage = 7
    yield stage, f"[{STAGE_NAMES[stage - 1]}] Resegmenting subtitles…", {}

    try:
        final_srt = _call_with_retry(
            resegment.resegment_srt,
            translated_srt,
            client=text_client, llm_model=LLM_MODEL_TEXT,
            usage_tracker=tracker,
        )
        Path(proj.final_srt).write_text(final_srt)
    except Exception as exc:
        yield stage, f"[{STAGE_NAMES[stage - 1]}] FAILED: {exc}", {"error": str(exc)}
        return

    yield stage, f"[{STAGE_NAMES[stage - 1]}] Done.", {}

    # -----------------------------------------------------------------------
    # Stage 8 — Synthesize (GPU)
    # -----------------------------------------------------------------------
    stage = 8
    yield stage, f"[{STAGE_NAMES[stage - 1]}] Synthesizing voice (first run downloads TTS model)…", {}

    srt_entries = parse_srt(proj.final_srt)

    _voice_theme = voice_theme if voice_mode == "theme" else None
    _clone_profile = voice_profile if voice_mode == "profile" else None
    _voice_sample = voice_sample if voice_mode == "clone" else None
    _voice_script = voice_script if voice_mode == "clone" else None

    # Auto-clone: extract a voice segment from source audio (CPU, no GPU needed)
    if voice_mode == "auto":
        yield stage, f"[{STAGE_NAMES[stage - 1]}] Auto-cloning voice from source audio…", {}
        try:
            auto_profile_dir = os.path.join(proj.root, "voice_profile")
            _voice_sample = profiles.create_auto_clone_profile(
                proj.audio, proj.source_srt, auto_profile_dir,
            )
        except Exception as exc:
            yield stage, f"[{STAGE_NAMES[stage - 1]}] Auto-clone failed: {exc} — falling back to narrator-m.", {}
            _voice_theme = "narrator-m"
            _voice_sample = None

    try:
        segment_info = _gpu_synthesize(
            tts_model_name="Qwen/Qwen3-TTS-12Hz-1.7B-Base",
            voice_sample=_voice_sample,
            voice_script=_voice_script,
            voice_theme=_voice_theme,
            clone_profile=_clone_profile,
            srt_entries=srt_entries,
            output_dir=proj.tts_segments_dir,
            target_language=target_language,
        )
    except TimeoutError:
        yield stage, f"[{STAGE_NAMES[stage - 1]}] GPU timeout — try a shorter clip.", {"error": "timeout"}
        return
    except Exception as exc:
        yield stage, f"[{STAGE_NAMES[stage - 1]}] FAILED: {exc}", {"error": str(exc)}
        return

    yield stage, f"[{STAGE_NAMES[stage - 1]}] Synthesis complete.", {}

    # -----------------------------------------------------------------------
    # Stage 9 — Assemble
    # -----------------------------------------------------------------------
    stage = 9
    yield stage, f"[{STAGE_NAMES[stage - 1]}] Assembling dubbed timeline…", {}

    try:
        original_duration = sf.info(proj.audio).duration
        assemble.assemble_timeline(
            segment_info, original_duration, proj.final_audio,
            tempo_mode=tempo_mode,
            max_tempo=max_tempo,
        )
        assemble.post_process(proj.final_audio, proj.audio, proj.final_audio)
    except Exception as exc:
        yield stage, f"[{STAGE_NAMES[stage - 1]}] FAILED: {exc}", {"error": str(exc)}
        return

    yield stage, f"[{STAGE_NAMES[stage - 1]}] Audio assembled.", {}

    # -----------------------------------------------------------------------
    # Stage 10 — Subtitle / Mux
    # -----------------------------------------------------------------------
    stage = 10
    yield stage, f"[{STAGE_NAMES[stage - 1]}] Creating final output…", {}

    try:
        if embed_subtitles and Path(proj.video).exists():
            font_file = None
            try:
                font_file = download_google_font(subtitle_font)
            except Exception:
                pass
            style = SubtitleStyle(
                font=subtitle_font,
                font_file=font_file,
                font_size=subtitle_font_size,
                font_color=subtitle_color,
                position=subtitle_position,
                bg_alpha=subtitle_bg_alpha,
                outline_width=subtitle_outline_width,
                bold=subtitle_bold,
                line_spacing=subtitle_line_spacing,
            )
            # Resolve subtitle source SRT
            if subtitle_source == "original":
                srt_for_burn = proj.source_srt
            else:
                srt_for_burn = proj.translated_raw_srt if os.path.exists(proj.translated_raw_srt) else proj.final_srt
            subtitle.burn_subtitles(proj.video, proj.final_video, srt_for_burn, style=style, audio_path=proj.final_audio)
            final_path = proj.final_video
        elif Path(proj.video).exists():
            assemble.mux_video(proj.video, proj.final_audio, proj.final_video)
            final_path = proj.final_video
        else:
            final_path = proj.final_audio
    except Exception as exc:
        yield stage, f"[{STAGE_NAMES[stage - 1]}] FAILED: {exc}", {"error": str(exc)}
        return

    # Copy to persistent temp dir
    result_file = shutil.copy2(final_path, os.path.join(result_tmp, os.path.basename(final_path)))
    result_srt = shutil.copy2(proj.final_srt, os.path.join(result_tmp, "subtitles.srt"))

    yield stage, f"[{STAGE_NAMES[stage - 1]}] Done.", {"final_path": result_file, "srt_path": result_srt}

    # Final sentinel
    yield len(STAGE_NAMES), "Pipeline complete.", {"final_path": result_file, "srt_path": result_srt, "done": True}