"""Subtitle cleanup: mask long-duration closing/greeting hallucinations.

Background (real run 20260905115050): after fixing the timing alignment,
subtitles still contain "closing/greeting hallucination words" - fixed
phrases like 'wan an / gan xie guan kan / gan xie nin de guan kan'
(good night / thanks for watching) that the ASR/LLM repeatedly emits on
empty segments, filling a full 30s block, unrelated to video content.
Some 2s 'good night' might be real dialogue, so it must be kept.

Plan (confirmed by user): after translation, mask subtitle entries whose
*display duration* exceeds a threshold AND whose text contains a greeting
hallucination token - replace the text with '-' so the downstream SRT/filter
pipeline drops it. The duration threshold protects short real greetings.

Threshold is derived from real run data: 30s hallucinations vs 2s real words,
a clear gap; default 15s (>=15s masks, <15s keeps).

Pure functions, unit-testable (tests/test_hallucination_mask.py).
"""

from __future__ import annotations

import re

# Greeting/closing hallucination tokens that LLM repeats on empty/end segments.
HALLUCINATION_TOKENS = (
    "谢谢观看", "感谢观看", "感谢收看", "谢谢收看", "感谢您的观看", "感谢您的收看",
    "晚安", "下次再见", "再会", "敬请期待", "感谢您的光临", "欢迎光临",
    "再见", "多谢观看", "观看愉快",
)

# Display-duration threshold (seconds): only mask entries longer than this.
DEFAULT_THRESHOLD_SECONDS = 15.0

_SRT_BLOCK = re.compile(
    r"(\d{2}:\d{2}:\d{2},\d{3})\s*-->\s*(\d{2}:\d{2}:\d{2},\d{3})\s*\n(.*?)(?=\n\s*\d+\s*\n|\Z)",
    re.DOTALL,
)


def mask_hallucination_text(
    entries: list[dict],
    threshold_seconds: float = DEFAULT_THRESHOLD_SECONDS,
) -> list[dict]:
    """Return a new list where long-duration greeting entries have text='-'.

    duration = end - start. Only entries whose duration >= threshold AND text
    contains any HALLUCINATION_TOKENS are masked. Input list is not mutated.
    """
    cleaned = []
    for entry in entries:
        duration = entry.get("end", 0.0) - entry.get("start", 0.0)
        text = entry.get("text", "")
        if duration >= threshold_seconds and any(t in text for t in HALLUCINATION_TOKENS):
            entry = dict(entry, text="-")
        cleaned.append(entry)
    return cleaned


def _ts_to_seconds(ts: str) -> float:
    """Convert an SRT timestamp HH:MM:SS,mmm to seconds (float)."""
    hours, minutes, rest = ts.split(":")
    seconds, millis = rest.split(",")
    return int(hours) * 3600 + int(minutes) * 60 + int(seconds) + int(millis) / 1000


def clean_srt_text(
    srt_text: str,
    threshold_seconds: float = DEFAULT_THRESHOLD_SECONDS,
) -> str:
    """Mask long-duration greeting hallucinations in an SRT string.

    Parses each cue's start/end/text, applies mask_hallucination_text, and
    rewrites the block keeping the original time line when not masked.
    """
    def _replace(match) -> str:
        start = _ts_to_seconds(match.group(1))
        end = _ts_to_seconds(match.group(2))
        text = match.group(3).strip()
        entry = {"start": start, "end": end, "text": text}
        cleaned = mask_hallucination_text([entry], threshold_seconds)
        new_text = cleaned[0]["text"]
        return f"{match.group(1)} --> {match.group(2)}\n{new_text}"

    return _SRT_BLOCK.sub(_replace, srt_text)