src/omnisub/scene_context.py

"""Trích ngữ cảnh đa phương thức cho từng cảnh: frame video + audio đoạn cảnh.

Dùng ffmpeg/ffprobe (gọi qua subprocess). Trên Colab L4 đã có sẵn ffmpeg.
Tất cả file tạm ghi vào `work_dir` để backend Omni nạp ở Bước 3 & 4.
"""

from __future__ import annotations

import shutil
import subprocess
from dataclasses import dataclass, field
from pathlib import Path
from typing import List, Optional

from .scenes import Scene


def _require(tool: str) -> str:
    path = shutil.which(tool)
    if not path:
        raise RuntimeError(
            f"Không tìm thấy `{tool}` trong PATH. Cần ffmpeg để trích frame/audio."
        )
    return path


def extract_audio(
    video: str | Path,
    out_path: str | Path,
    *,
    start: Optional[float] = None,
    end: Optional[float] = None,
    sample_rate: int = 16000,
) -> Path:
    """Trích audio mono 16 kHz (mặc định) cho [start, end]. Trả về path WAV."""
    ffmpeg = _require("ffmpeg")
    out = Path(out_path)
    out.parent.mkdir(parents=True, exist_ok=True)
    cmd = [ffmpeg, "-y", "-loglevel", "error"]
    if start is not None:
        cmd += ["-ss", f"{start:.3f}"]
    if end is not None and start is not None:
        cmd += ["-t", f"{max(0.0, end - start):.3f}"]
    cmd += [
        "-i", str(video),
        "-vn",
        "-ac", "1",
        "-ar", str(sample_rate),
        "-c:a", "pcm_s16le",
        str(out),
    ]
    subprocess.run(cmd, check=True)
    return out


def extract_frame(
    video: str | Path,
    out_path: str | Path,
    *,
    timestamp: float,
    width: Optional[int] = 768,
) -> Path:
    """Trích 1 frame tại `timestamp` (giây). Trả về path ảnh JPG."""
    ffmpeg = _require("ffmpeg")
    out = Path(out_path)
    out.parent.mkdir(parents=True, exist_ok=True)
    vf = f"scale={width}:-2" if width else "scale=iw:ih"
    cmd = [
        ffmpeg, "-y", "-loglevel", "error",
        "-ss", f"{timestamp:.3f}",
        "-i", str(video),
        "-frames:v", "1",
        "-vf", vf,
        str(out),
    ]
    subprocess.run(cmd, check=True)
    return out


@dataclass
class SceneContext:
    """Tài nguyên đa phương thức đã trích cho một cảnh."""

    scene_id: int
    audio_path: Optional[Path] = None
    frame_paths: List[Path] = field(default_factory=list)


def build_scene_context(
    video: str | Path,
    scene: Scene,
    work_dir: str | Path,
    *,
    n_frames: int = 2,
    sample_rate: int = 16000,
    frame_width: int = 768,
) -> SceneContext:
    """Trích audio toàn cảnh + `n_frames` frame rải đều trong cảnh."""
    work = Path(work_dir)
    work.mkdir(parents=True, exist_ok=True)

    audio_path = extract_audio(
        video,
        work / f"scene_{scene.scene_id:04d}.wav",
        start=scene.start,
        end=scene.end,
        sample_rate=sample_rate,
    )

    frames: List[Path] = []
    dur = max(0.001, scene.duration)
    for i in range(max(1, n_frames)):
        # rải đều, tránh đúng biên cắt cảnh
        frac = (i + 1) / (n_frames + 1)
        ts = scene.start + frac * dur
        frames.append(
            extract_frame(
                video,
                work / f"scene_{scene.scene_id:04d}_f{i}.jpg",
                timestamp=ts,
                width=frame_width,
            )
        )

    return SceneContext(scene_id=scene.scene_id, audio_path=audio_path, frame_paths=frames)


def extract_speaker_clips(
    video: str | Path,
    turns: List["SpeakerTurn"],  # type: ignore[name-defined]
    work_dir: str | Path,
    *,
    clips_per_speaker: int = 5,
    max_seconds_per_clip: float = 8.0,
    sample_rate: int = 16000,
) -> dict[str, List[Path]]:
    """Trích các clip audio đại diện cho từng speaker (cho Bước 3 — hồ sơ giọng).

    Chọn các turn dài nhất của mỗi speaker, cắt ≤ `max_seconds_per_clip`.
    """
    from collections import defaultdict

    work = Path(work_dir)
    work.mkdir(parents=True, exist_ok=True)

    by_speaker: dict[str, list] = defaultdict(list)
    for t in turns:
        by_speaker[t.speaker].append(t)

    result: dict[str, List[Path]] = {}
    for speaker, sp_turns in by_speaker.items():
        sp_turns = sorted(sp_turns, key=lambda x: x.duration, reverse=True)
        clips: List[Path] = []
        for i, t in enumerate(sp_turns[:clips_per_speaker]):
            end = min(t.end, t.start + max_seconds_per_clip)
            clips.append(
                extract_audio(
                    video,
                    work / f"{speaker}_clip{i}.wav",
                    start=t.start,
                    end=end,
                    sample_rate=sample_rate,
                )
            )
        result[speaker] = clips
    return result