#!/usr/bin/env python3
"""把 fastasr transcribe_long 的带时间戳文本转成 SRT / VTT 字幕。

输入：每行形如  [开始,结束]  文本   （开始/结束为 fastasr 的时间戳，如 0:4.640 或 1:23:45.120）
用法：
  python3 to_subtitle.py raw.txt          > out.srt
  python3 to_subtitle.py raw.txt vtt       > out.vtt
  cat raw.txt | python3 to_subtitle.py - vtt > out.vtt   # 也支持 stdin
仅用标准库，无第三方依赖。
"""
import re
import sys

SEG_RE = re.compile(r"\[\s*([\d:.]+)\s*,\s*([\d:.]+)\s*\]\s*(.*)")


def parse_ts(s):
    """'0:4.640' / '1:23:45.120' / '5.0' -> 秒(float)。最后一段是秒，往前依次是分、时。"""
    parts = s.split(":")
    sec = float(parts[-1])
    if len(parts) >= 2:
        sec += int(parts[-2]) * 60
    if len(parts) >= 3:
        sec += int(parts[-3]) * 3600
    return sec


def fmt(sec, vtt=False):
    h = int(sec // 3600)
    m = int((sec % 3600) // 60)
    whole = int(sec % 60)
    ms = int(round((sec - int(sec)) * 1000))
    if ms == 1000:  # 进位
        ms = 0
        whole += 1
    sep = "." if vtt else ","
    return f"{h:02d}:{m:02d}:{whole:02d}{sep}{ms:03d}"


def main():
    if len(sys.argv) < 2:
        sys.exit("用法: to_subtitle.py <raw.txt|-> [vtt]")
    src = sys.argv[1]
    vtt = len(sys.argv) > 2 and sys.argv[2].lower().lstrip("-") == "vtt"
    raw = sys.stdin.read() if src == "-" else open(src, encoding="utf-8").read()

    cues = []
    for line in raw.splitlines():
        line = line.strip()
        if not line:
            continue
        m = SEG_RE.match(line)
        if not m:
            continue
        text = m.group(3).strip()
        if not text:
            continue
        cues.append((parse_ts(m.group(1)), parse_ts(m.group(2)), text))

    if not cues:
        sys.exit("未解析到任何带时间戳的片段。请确认转写时 plain_text 为 false（保留时间戳）。")

    out = []
    if vtt:
        out.append("WEBVTT\n")
    for i, (st, en, text) in enumerate(cues, 1):
        if not vtt:
            out.append(str(i))
        out.append(f"{fmt(st, vtt)} --> {fmt(en, vtt)}")
        out.append(text)
        out.append("")
    sys.stdout.write("\n".join(out).rstrip() + "\n")


if __name__ == "__main__":
    main()