#!/usr/bin/env python3
"""pgs2ref.py — 把蓝光碟内的 PGS 字幕转成"仅时间轴"的 SRT，当作 alass 的 reference

PGS(Presentation Graphic Stream) 是图形字幕，没有文本。但 alass 的参考字幕模式
只比对时间轴分布、不看文字内容，所以不需要 OCR，直接按每段 PCS 的 PTS 造一条
占位字幕即可，精度与 OCR 版完全一致，速度快上百倍。

用法:
  ./pgs2ref.py eng.sup -o ref.srt                    # 直接吃 .sup
  ./pgs2ref.py movie.mkv --list                      # 列出容器内所有字幕流
  ./pgs2ref.py movie.mkv --stream 2 -o ref.srt       # 从容器里抠第 3 条字幕流
  ./pgs2ref.py movie.mkv --lang eng -o ref.srt       # 按语言挑
  ./pgs2ref.py 00001.m2ts --lang eng -o ref.srt      # 直接吃 m2ts 也行
"""

import argparse
import json
import subprocess
import sys
import tempfile
from pathlib import Path

PCS = 0x16
# 同一字幕的"更新帧"间隔阈值(秒): 多窗口/多次刷新的 PCS 会被合并成一条
MERGE_GAP = 0.15
DEFAULT_DUR = 2.0


def parse_sup(path):
    """解析 .sup, 返回 [(start, end), ...] 秒"""
    data = Path(path).read_bytes()
    n = len(data)
    i = 0
    events = []  # (pts, num_objects)
    while i + 13 <= n:
        if data[i:i + 2] != b"PG":
            j = data.find(b"PG", i + 1)
            if j < 0:
                break
            i = j
            continue
        pts = int.from_bytes(data[i + 2:i + 6], "big")
        seg_type = data[i + 10]
        seg_size = int.from_bytes(data[i + 11:i + 13], "big")
        payload = data[i + 13:i + 13 + seg_size]
        if seg_type == PCS and len(payload) >= 11:
            events.append((pts / 90000.0, payload[10]))
        i += 13 + seg_size

    cues = []
    cur = None
    for pts, nobj in events:
        if nobj == 0:                      # clear: 当前字幕结束
            if cur is not None:
                cur[1] = pts
                cues.append(cur)
                cur = None
            continue
        if cur is not None and pts - cur[0] < MERGE_GAP:
            continue                        # 同一条字幕的刷新
        if cur is not None:                 # 上一条没收到 clear, 用新条目起点兜底
            cur[1] = pts
            cues.append(cur)
        cur = [pts, None]
    if cur is not None:
        cues.append([cur[0], cur[0] + DEFAULT_DUR])

    out = []
    for s, e in cues:
        if e is None or e <= s:
            e = s + DEFAULT_DUR
        out.append((s, e))
    return out


def fmt_ts(t):
    t = max(t, 0.0)
    ms = int(round(t * 1000))
    h, ms = divmod(ms, 3_600_000)
    m, ms = divmod(ms, 60_000)
    s, ms = divmod(ms, 1000)
    return f"{h:02d}:{m:02d}:{s:02d},{ms:03d}"


def probe_streams(path):
    cmd = ["ffprobe", "-v", "error", "-select_streams", "s",
           "-show_entries", "stream=index,codec_name:stream_tags=language,title",
           "-of", "json", str(path)]
    r = subprocess.run(cmd, capture_output=True, text=True)
    if r.returncode != 0:
        sys.exit(f"ffprobe 失败: {r.stderr.strip()}")
    return json.loads(r.stdout or "{}").get("streams", [])


def pick_stream(streams, lang, idx):
    if idx is not None:
        if idx >= len(streams):
            sys.exit(f"--stream {idx} 越界, 共 {len(streams)} 条字幕流")
        return streams[idx]
    if lang:
        for s in streams:
            if s.get("tags", {}).get("language", "").lower() == lang.lower():
                return s
        sys.exit(f"没有语言为 {lang} 的字幕流, 用 --list 看实际标签")
    for s in streams:
        if s.get("codec_name") == "hdmv_pgs_subtitle":
            return s
    if streams:
        return streams[0]
    sys.exit("容器内没有字幕流")


def extract_sup(path, stream_index, tmp):
    # -copyts: 保持原始时间戳, 否则 ffmpeg 会把首个包归零, 让整条字幕整体偏移
    cmd = ["ffmpeg", "-v", "error", "-y", "-copyts", "-i", str(path),
           "-map", f"0:{stream_index}", "-c", "copy", "-f", "sup", str(tmp)]
    r = subprocess.run(cmd, capture_output=True, text=True)
    if r.returncode != 0:
        sys.exit(f"ffmpeg 提取失败: {r.stderr.strip()}")
    return tmp


def main():
    p = argparse.ArgumentParser(description="PGS 字幕 → 仅时间轴的 reference SRT")
    p.add_argument("input", help=".sup / .m2ts / .mkv / .ts")
    p.add_argument("-o", "--out", help="输出 SRT, 默认 <原名>.ref.srt")
    p.add_argument("--list", action="store_true", help="只列出字幕流")
    p.add_argument("--stream", type=int, help="第 N 条字幕流(从 0 开始)")
    p.add_argument("--lang", help="按语言挑, 如 eng / chi")
    p.add_argument("--text", default="x", help="占位文本, 默认 'x'")
    args = p.parse_args()

    src = Path(args.input)
    if not src.exists():
        sys.exit(f"找不到文件: {src}")

    tmp = None
    if src.suffix.lower() == ".sup":
        sup_path = src
    else:
        streams = probe_streams(src)
        if args.list:
            if not streams:
                print("没有字幕流")
                return
            for i, s in enumerate(streams):
                tags = s.get("tags", {})
                print(f"[{i}] {s.get('codec_name','?'):<22} "
                      f"lang={tags.get('language','-'):<6} title={tags.get('title','-')}")
            return
        if not streams:
            sys.exit("容器内没有字幕流")
        st = pick_stream(streams, args.lang, args.stream)
        fd = tempfile.NamedTemporaryFile(suffix=".sup", delete=False)
        fd.close()
        tmp = Path(fd.name)
        extract_sup(src, st["index"], tmp)
        sup_path = tmp
        print(f"已提取字幕流 [{st['index']}] {st.get('codec_name')}", file=sys.stderr)

    cues = parse_sup(sup_path)
    if tmp:
        tmp.unlink(missing_ok=True)
    if not cues:
        sys.exit("没解析到任何 PGS 条目, 确认这条流是 hdmv_pgs_subtitle")

    out = Path(args.out) if args.out else src.with_suffix(".ref.srt")
    blocks = [f"{i}\n{fmt_ts(s)} --> {fmt_ts(e)}\n{args.text}"
              for i, (s, e) in enumerate(cues, 1)]
    out.write_text("\n\n".join(blocks) + "\n", encoding="utf-8")

    total = cues[-1][1]
    print(f"{len(cues)} 条, 首条 {fmt_ts(cues[0][0])}, 末条 {fmt_ts(total)} → {out}")
    if cues[0][0] < 5:
        print("警告: 首条字幕出现在 5 秒内。若正片开头本该无字幕, 说明中转环节"
              "把时间戳归零了(常见于 ffmpeg/mkvmerge 重封装), reference 会整体偏小。"
              "直接用 MakeMKV 的输出或从 m2ts 原样提取可避免。")
    print(f"下一步: alass {out} 中字.srt 中字_fixed.srt --split-penalty 10")


if __name__ == "__main__":
    main()