src/omnisub/srt.py

"""I/O SRT: dataclass `Cue`, parse/write, ngân sách ký tự theo thời lượng.

Port sạch từ bản cũ — không phụ thuộc gì ngoài stdlib để chạy được local ở Bước 1.
"""

from __future__ import annotations

import re
from dataclasses import dataclass, field
from pathlib import Path
from typing import Iterable, List, Optional

# 00:01:02,500 --> 00:01:05,000
_TIME_RE = re.compile(
    r"(?P<h>\d{1,2}):(?P<m>\d{2}):(?P<s>\d{2})[.,](?P<ms>\d{1,3})"
)
_ARROW_RE = re.compile(r"-->")


@dataclass
class Cue:
    """Một dòng phụ đề."""

    index: int
    start: float           # giây
    end: float             # giây
    text: str              # văn bản gốc (ZH)
    speaker: Optional[str] = None   # gán ở Bước 2 (diarization)
    translation: Optional[str] = None  # bản dịch VN (Bước 4)
    note: Optional[str] = None         # ghi chú xưng hô của model
    scene_id: Optional[int] = None     # gán ở Bước 1 (gom cảnh)

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

    @property
    def out_text(self) -> str:
        """Văn bản dùng khi ghi ra: ưu tiên bản dịch, fallback về gốc."""
        return self.translation if self.translation else self.text


def parse_time(token: str) -> float:
    """'00:01:02,500' -> giây (float)."""
    m = _TIME_RE.search(token)
    if not m:
        raise ValueError(f"Không đọc được mốc thời gian: {token!r}")
    h = int(m.group("h"))
    mi = int(m.group("m"))
    s = int(m.group("s"))
    ms = int(m.group("ms").ljust(3, "0"))  # '5' -> 500
    return h * 3600 + mi * 60 + s + ms / 1000.0


def seconds_to_srt_time(seconds: float) -> str:
    """giây -> '00:01:02,500'."""
    if seconds < 0:
        seconds = 0.0
    total_ms = int(round(seconds * 1000.0))
    ms = total_ms % 1000
    total_s = total_ms // 1000
    s = total_s % 60
    total_m = total_s // 60
    m = total_m % 60
    h = total_m // 60
    return f"{h:02d}:{m:02d}:{s:02d},{ms:03d}"


def _normalize_text(raw_lines: List[str]) -> str:
    """Gộp các dòng văn bản của một cue, bỏ thẻ định dạng đơn giản."""
    text = "\n".join(line.rstrip() for line in raw_lines).strip()
    return text


def parse_srt(path: str | Path, encoding: str = "utf-8-sig") -> List[Cue]:
    """Đọc file SRT -> danh sách Cue. Bền với index lệch, dòng trống dư."""
    p = Path(path)
    content = p.read_text(encoding=encoding, errors="replace")
    return parse_srt_text(content)


def parse_srt_text(content: str) -> List[Cue]:
    content = content.replace("\r\n", "\n").replace("\r", "\n")
    blocks = re.split(r"\n\s*\n", content.strip())
    cues: List[Cue] = []
    auto_index = 0
    for block in blocks:
        lines = [ln for ln in block.split("\n")]
        if not lines:
            continue
        # bỏ dòng trống đầu
        while lines and not lines[0].strip():
            lines.pop(0)
        if not lines:
            continue

        idx_line = lines[0].strip()
        time_line_pos = 0
        explicit_index: Optional[int] = None
        if idx_line.isdigit() and len(lines) >= 2 and _ARROW_RE.search(lines[1]):
            explicit_index = int(idx_line)
            time_line_pos = 1
        elif _ARROW_RE.search(lines[0]):
            time_line_pos = 0
        else:
            # block không hợp lệ (vd metadata) -> bỏ qua
            continue

        time_line = lines[time_line_pos]
        parts = _ARROW_RE.split(time_line)
        if len(parts) != 2:
            continue
        try:
            start = parse_time(parts[0])
            end = parse_time(parts[1])
        except ValueError:
            continue

        text = _normalize_text(lines[time_line_pos + 1 :])
        auto_index += 1
        cues.append(
            Cue(
                index=explicit_index if explicit_index is not None else auto_index,
                start=start,
                end=end,
                text=text,
            )
        )
    return cues


def write_srt(cues: Iterable[Cue], path: str | Path, encoding: str = "utf-8") -> None:
    """Ghi danh sách Cue ra file SRT (dùng `out_text`)."""
    p = Path(path)
    p.parent.mkdir(parents=True, exist_ok=True)
    p.write_text(render_srt(cues), encoding=encoding)


def render_srt(cues: Iterable[Cue]) -> str:
    chunks: List[str] = []
    for i, cue in enumerate(cues, start=1):
        chunks.append(
            f"{i}\n"
            f"{seconds_to_srt_time(cue.start)} --> {seconds_to_srt_time(cue.end)}\n"
            f"{cue.out_text}".rstrip()
        )
    return "\n\n".join(chunks) + "\n"


def subtitle_char_budget(
    duration: float,
    chars_per_sec: float = 22.0,
    max_line_chars: int = 52,
    max_lines: int = 2,
) -> int:
    """Số ký tự tối đa nên hiển thị cho một cue, theo thời lượng và giới hạn dòng.

    Lấy min giữa (thời lượng × tốc độ đọc) và (số ký tự tối đa của các dòng).
    """
    by_time = int(duration * chars_per_sec)
    by_lines = max_line_chars * max_lines
    # tối thiểu cho 1 dòng ngắn để không trả về 0 với cue siêu ngắn
    return max(by_lines if by_time <= 0 else min(by_time, by_lines), max_line_chars)