src/omnisub/scenes.py

"""Gom cảnh và ghép cue vụn (sentence-aware, same-speaker).

`group_scenes`: cắt phim thành các `Scene` theo khoảng lặng và giới hạn độ dài.
`merge_adjacent_pairs`: ghép cue ngắn liền kề (cùng người nói, chưa kết câu).
"""

from __future__ import annotations

from dataclasses import dataclass, field
from typing import List, Optional

from .srt import Cue

# Dấu kết thúc câu (CJK + Latin) — dùng để biết một cue đã "trọn câu" chưa.
_SENTENCE_ENDERS = "。！？!?…」』”》"
_CONTINUATION_ENDERS = "，、,;；:：—-（(【「『《“"


@dataclass
class Scene:
    """Một cảnh: cụm cue gần nhau về thời gian."""

    scene_id: int
    cues: List[Cue] = field(default_factory=list)

    @property
    def start(self) -> float:
        return self.cues[0].start if self.cues else 0.0

    @property
    def end(self) -> float:
        return self.cues[-1].end if self.cues else 0.0

    @property
    def duration(self) -> float:
        return max(0.0, self.end - self.start)

    @property
    def text(self) -> str:
        return "\n".join(c.text for c in self.cues)

    @property
    def speakers(self) -> List[str]:
        seen: List[str] = []
        for c in self.cues:
            if c.speaker and c.speaker not in seen:
                seen.append(c.speaker)
        return seen


def group_scenes(
    cues: List[Cue],
    max_gap: float = 1.5,
    max_cues: int = 4,
    max_dur: float = 20.0,
) -> List[Scene]:
    """Gom cue thành cảnh.

    Cắt cảnh mới khi: khoảng lặng > `max_gap`, hoặc cảnh đã đủ `max_cues`,
    hoặc vượt `max_dur` giây.
    """
    scenes: List[Scene] = []
    if not cues:
        return scenes

    current: List[Cue] = []
    scene_id = 0
    for cue in cues:
        if not current:
            current = [cue]
            continue

        gap = cue.start - current[-1].end
        dur = cue.end - current[0].start
        if gap > max_gap or len(current) >= max_cues or dur > max_dur:
            scene_id += 1
            scenes.append(_finalize_scene(scene_id, current))
            current = [cue]
        else:
            current.append(cue)

    if current:
        scene_id += 1
        scenes.append(_finalize_scene(scene_id, current))
    return scenes


def _finalize_scene(scene_id: int, cues: List[Cue]) -> Scene:
    for c in cues:
        c.scene_id = scene_id
    return Scene(scene_id=scene_id, cues=list(cues))


def _ends_sentence(text: str) -> bool:
    t = text.rstrip()
    if not t:
        return False
    return t[-1] in _SENTENCE_ENDERS


def _ends_continuation(text: str) -> bool:
    t = text.rstrip()
    if not t:
        return False
    return t[-1] in _CONTINUATION_ENDERS


def merge_adjacent_pairs(
    cues: List[Cue],
    max_gap: float = 0.4,
    max_merged_chars: int = 40,
) -> List[Cue]:
    """Ghép cue vụn liền kề thành câu trọn vẹn hơn.

    Điều kiện ghép cue[i] với cue[i+1]:
      - khoảng cách ≤ `max_gap`,
      - cùng speaker (hoặc cả hai chưa gán),
      - cue[i] CHƯA kết thúc câu,
      - tổng độ dài sau ghép ≤ `max_merged_chars`.
    """
    if not cues:
        return []

    merged: List[Cue] = []
    buf = _clone(cues[0])
    for nxt in cues[1:]:
        gap = nxt.start - buf.end
        # Cùng speaker đã gán → ghép thoải mái. Khi cả hai chưa gán (chưa
        # diarize) chỉ ghép nếu câu trước kết thúc bằng dấu nối, tránh gộp
        # nhầm lời của hai người khác nhau.
        both_unknown = buf.speaker is None and nxt.speaker is None
        same_known_speaker = buf.speaker is not None and buf.speaker == nxt.speaker
        can_merge_speaker = same_known_speaker or (
            both_unknown and _ends_continuation(buf.text)
        )
        combined_len = len(buf.text.replace("\n", "")) + len(nxt.text.replace("\n", ""))
        if (
            gap <= max_gap
            and can_merge_speaker
            and not _ends_sentence(buf.text)
            and combined_len <= max_merged_chars
        ):
            buf.end = nxt.end
            sep = "" if buf.text and buf.text[-1] in _CONTINUATION_ENDERS else " "
            buf.text = (buf.text + sep + nxt.text).strip()
        else:
            merged.append(buf)
            buf = _clone(nxt)
    merged.append(buf)

    for i, c in enumerate(merged, start=1):
        c.index = i
    return merged


def _clone(cue: Cue) -> Cue:
    return Cue(
        index=cue.index,
        start=cue.start,
        end=cue.end,
        text=cue.text,
        speaker=cue.speaker,
        translation=cue.translation,
        note=cue.note,
        scene_id=cue.scene_id,
    )