renderer/subtitles/generator.py

from __future__ import annotations

import html
from dataclasses import dataclass
from pathlib import Path

from renderer.templates import get_template


@dataclass
class SubtitleEvent:
    start: float
    end: float
    text: str


class SubtitleGenerator:
    def from_scenes(self, scenes: list, total_duration: float | None = None) -> list[SubtitleEvent]:
        events: list[SubtitleEvent] = []
        for scene in scenes:
            if scene.caption:
                events.append(SubtitleEvent(scene.start, scene.start + scene.duration, scene.caption))
        if not events and total_duration:
            events.append(SubtitleEvent(0, total_duration, ""))
        return events

    def write_srt(self, events: list[SubtitleEvent], output: Path) -> Path:
        lines: list[str] = []
        for idx, event in enumerate(events, start=1):
            lines.extend([str(idx), f"{_srt_time(event.start)} --> {_srt_time(event.end)}", event.text, ""])
        output.write_text("\n".join(lines), encoding="utf-8")
        return output

    def write_ass(self, events: list[SubtitleEvent], output: Path, template_key: str) -> Path:
        template = get_template(template_key)
        body = [
            "[Script Info]",
            "ScriptType: v4.00+",
            "PlayResX: 1080",
            "PlayResY: 1920",
            "",
            "[V4+ Styles]",
            "Format: Name,Fontname,Fontsize,PrimaryColour,SecondaryColour,OutlineColour,BackColour,"
            "Bold,Italic,Underline,StrikeOut,ScaleX,ScaleY,Spacing,Angle,BorderStyle,Outline,Shadow,"
            "Alignment,MarginL,MarginR,MarginV,Encoding",
            template.ass_style(),
            "",
            "[Events]",
            "Format: Layer,Start,End,Style,Name,MarginL,MarginR,MarginV,Effect,Text",
        ]
        for event in events:
            text = _ass_escape(event.text)
            if template.effect == "karaoke":
                text = _karaoke_text(text, event.end - event.start)
            elif template.effect == "zoom":
                text = r"{\t(0,180,\fscx115\fscy115)\t(180,360,\fscx100\fscy100)}" + text
            elif template.effect == "bounce":
                text = r"{\t(0,120,\frz-2)\t(120,240,\frz2)\t(240,360,\frz0)}" + text
            body.append(f"Dialogue: 0,{_ass_time(event.start)},{_ass_time(event.end)},Default,,0,0,0,,{text}")
        output.write_text("\n".join(body), encoding="utf-8")
        return output


def _karaoke_text(text: str, duration: float) -> str:
    words = text.split()
    if not words:
        return text
    centiseconds = max(1, int(duration * 100 / len(words)))
    return "".join(f"{{\\k{centiseconds}}}{word} " for word in words).strip()


def _srt_time(seconds: float) -> str:
    ms = int(round(seconds * 1000))
    h, rem = divmod(ms, 3600000)
    m, rem = divmod(rem, 60000)
    s, ms = divmod(rem, 1000)
    return f"{h:02d}:{m:02d}:{s:02d},{ms:03d}"


def _ass_time(seconds: float) -> str:
    cs = int(round(seconds * 100))
    h, rem = divmod(cs, 360000)
    m, rem = divmod(rem, 6000)
    s, cs = divmod(rem, 100)
    return f"{h}:{m:02d}:{s:02d}.{cs:02d}"


def _ass_escape(text: str) -> str:
    return html.escape(text).replace("\n", r"\N").replace("{", r"\{").replace("}", r"\}")