subtitle_parser.py

"""
subtitle_parser.py — Subtitle Parsing & Writing
Handles .srt / .vtt / .ass via pysubs2.
"""

import copy
import logging
import os
import re
import shutil
import tempfile

import pysubs2

logger = logging.getLogger(__name__)

WORKSPACE_DIR = os.path.join(tempfile.gettempdir(), "subsync_workspace")
os.makedirs(WORKSPACE_DIR, exist_ok=True)

# ASS styles that should NOT be translated (romaji, karaoke, etc.)
SKIP_STYLES = {
    "EDR", "Romaji", "JP", "Karaoke", "KAR",
    "Signs", "Lyrics_JP", "Romaji2", "OP_Romaji", "ED_Romaji",
}

_EXT_MAP = {".srt": "srt", ".vtt": "vtt", ".ass": "ass", ".ssa": "ass"}

_ASS_TAG_RE      = re.compile(r"\{[^}]*\}")
_HTML_TAG_RE     = re.compile(r"<[^>]+>")
_LEADING_TAGS_RE = re.compile(r"^(\{[^}]*\})+")


def is_skip_style(style_name: str) -> bool:
    """Return True if this ASS style should be skipped (not translated)."""
    if not style_name:
        return False
    s = style_name.strip()
    if s in SKIP_STYLES:
        return True
    sl = s.lower()
    return any(k in sl for k in ["romaji", "karaoke", "kar_", "_kar", "furigana"])


def _detect_format(filename: str) -> str:
    ext = os.path.splitext(filename.lower())[1]
    fmt = _EXT_MAP.get(ext)
    if not fmt:
        raise ValueError(
            f"'{ext}' format ကို support မလုပ်ပါ — .srt / .vtt / .ass သာ လက်ခံသည်"
        )
    return fmt


def _strip_tags(text: str) -> str:
    text = _ASS_TAG_RE.sub("", text)
    text = _HTML_TAG_RE.sub("", text)
    return text.strip()


def parse_subtitle(file_bytes: bytes, filename: str):
    """Parse subtitle bytes → (SSAFile, fmt_string)."""
    fmt = _detect_format(filename)

    text = None
    for enc in ("utf-8-sig", "utf-8", "latin-1", "cp1252"):
        try:
            text = file_bytes.decode(enc)
            break
        except (UnicodeDecodeError, LookupError):
            continue
    if text is None:
        raise ValueError("ဖိုင် encoding ဖတ်၍ မရပါ")

    text = text.lstrip("\ufeff")

    try:
        subs = pysubs2.SSAFile.from_string(text)
        logger.info("Parsed %d events from '%s' (%s)", len(subs), filename, fmt)
        return subs, fmt
    except Exception as e1:
        logger.warning("from_string failed: %s", e1)

    try:
        suffix = os.path.splitext(filename)[1] or ("." + fmt)
        fd, tmp = tempfile.mkstemp(suffix=suffix)
        try:
            os.write(fd, file_bytes)
        finally:
            os.close(fd)
        subs = pysubs2.SSAFile.load(tmp)
        os.unlink(tmp)
        return subs, fmt
    except Exception as e2:
        raise ValueError(
            f"Parse မရပါ:\nAttempt 1: {e1}\nAttempt 2: {e2}"
        ) from e2


def get_plain_lines(ssafile) -> list:
    """
    Return list[str] same length as ssafile.events.
    - Comment events  → "" (keep original, not displayed)
    - Skip styles     → "" (EDR/romaji/karaoke etc.)
    - Normal events   → clean text (tags stripped)
    Empty string → replace_lines keeps original text unchanged.
    """
    lines = []
    for event in ssafile.events:
        # Skip ASS Comment lines (editor notes, not shown to viewer)
        if getattr(event, "is_comment", False):
            lines.append("")
            continue

        # Skip romaji / karaoke / non-English style tracks
        if is_skip_style(event.style):
            lines.append("")
            continue

        raw   = event.text.replace("\\N", "\n").replace("\\n", "\n")
        clean = _strip_tags(raw)
        lines.append(clean)
    return lines


def replace_lines(ssafile, new_lines: list):
    """
    Return new SSAFile with translated text. Timestamps never touched.
    Leading ASS override tags {\\pos(...)\\be5} preserved from original.
    Empty translation → keep original event.text.
    """
    n_events = len(ssafile.events)
    n_lines  = len(new_lines)

    if n_lines != n_events:
        raise ValueError(
            f"TIMESTAMP SAFETY: translated lines ({n_lines}) "
            f"!= subtitle events ({n_events})"
        )

    result = copy.deepcopy(ssafile)

    for i, event in enumerate(result.events):
        translated = str(new_lines[i]).strip()

        # Empty → keep original (skip-style or fallback)
        if not translated:
            continue

        # Re-prepend leading inline tags from original
        original_text = ssafile.events[i].text
        leading_m     = _LEADING_TAGS_RE.match(original_text)
        tag_prefix    = leading_m.group(0) if leading_m else ""

        event.text = tag_prefix + translated.replace("\n", "\\N")

    return result


def write_subtitle(ssafile, fmt: str) -> bytes:
    try:
        return ssafile.to_string(fmt).encode("utf-8")
    except Exception as exc:
        raise RuntimeError(f"Subtitle write မရပါ: {exc}") from exc


def output_filename(original_name: str) -> str:
    base, ext = os.path.splitext(original_name)
    return f"{base}.my{ext}"


def clean_workspace() -> None:
    try:
        if os.path.exists(WORKSPACE_DIR):
            shutil.rmtree(WORKSPACE_DIR)
        os.makedirs(WORKSPACE_DIR, exist_ok=True)
    except Exception as exc:
        raise RuntimeError(f"Workspace ရှင်းမရပါ: {exc}") from exc


def workspace_size_mb() -> float:
    total = 0
    for dirpath, _, filenames in os.walk(WORKSPACE_DIR):
        for f in filenames:
            try:
                total += os.path.getsize(os.path.join(dirpath, f))
            except OSError:
                pass
    return round(total / (1024 * 1024), 2)