#!/usr/bin/env python3
"""
Script to SRT with Timing · 剧本到 SRT 时间轴生成器
==========================​=====================
用剧本台词和配音时长生成正式 SRT，别手写测试时间轴。

功能：
  1. 读取剧本台词（JSON/MD/TXT）
  2. 读取配音文件（获取每条台词的实际时长）
  3. 自动计算时间轴（考虑角色对话间隔）
  4. 生成标准 SRT 字幕文件
  5. 支持多角色、旁白、音效标注

依赖：
  pip install pysrt

输入格式：
  # 剧本 JSON 格式（推荐）
  [
    {"id": 1, "character": "苏白", "text": "付费才能修仙？", "audio_file": "./audio/subai_001.wav"},
    {"id": 2, "character": "旁白", "text": "天道宗门口", "audio_file": "./audio/narrator_001.wav"}
  ]

  # 剧本 MD 格式（兼容）
  ## E1-SHOT01
  **苏白**：付费才能修仙？（audio: subai_001.wav）
  **旁白**：天道宗门口（audio: narrator_001.wav）

输出 SRT 格式：
  1
  00:00:01,000 --> 00:00:03,500
  付费才能修仙？

  2
  00:00:03,800 --> 00:00:05,200
  天道宗门口

用法：
  # 从剧本 JSON + 配音文件生成 SRT
  python script-to-srt-with-timing.py --script script.json --output subtitles.srt

  # 指定配音目录（自动匹配 audio_file）
  python script-to-srt-with-timing.py --script script.json --audio-dir ./audio/ --output subtitles.srt

  # 从 MD 剧本生成
  python script-to-srt-with-timing.py --script script.md --output subtitles.srt

  # 自定义对话间隔（秒）
  python script-to-srt-with-timing.py --script script.json --gap 0.3 --output subtitles.srt

路径：
  video-ai-system/engines/subtitle-pipeline/srt-tools/script-to-srt-with-timing.py
"""

import argparse
import json
import os
import sys
import re
from pathlib import Path

try:
    import pysrt
except ImportError:
    print("[ERROR] 缺少依赖：pysrt")
    print("请先安装：pip install pysrt")
    sys.exit(1)


def parse_script_json(script_path: str) -> list:
    """
    解析剧本 JSON 文件
    
    :param script_path: 剧本 JSON 文件路径
    :return: 台词列表 [{"id", "character", "text", "audio_file", "start", "end"}]
    """
    with open(script_path, "r", encoding="utf-8") as f:
        script = json.load(f)
    
    print(f"[INFO] 读取剧本 JSON：{len(script)} 条台词")
    return script


def parse_script_md(script_path: str) -> list:
    """
    解析剧本 MD 文件（兼容格式）
    
    :param script_path: 剧本 MD 文件路径
    :return: 台词列表
    """
    with open(script_path, "r", encoding="utf-8") as f:
        lines = f.readlines()
    
    script = []
    current_id = 0
    
    for line in lines:
        # 匹配：**角色**：台词（audio: file.wav）
        match = re.match(r"\*\*(.+?)\*\*[：:]\s*(.+?)\s*\(audio:\s*(.+?)\)", line.strip())
        if match:
            character = match.group(1).strip()
            text = match.group(2).strip()
            audio_file = match.group(3).strip()
            
            script.append({
                "id": current_id + 1,
                "character": character,
                "text": text,
                "audio_file": audio_file
            })
            current_id += 1
    
    print(f"[INFO] 读取剧本 MD：{len(script)} 条台词")
    return script


def get_audio_duration(audio_path: str) -> float:
    """
    获取音频文件时长（秒）
    
    :param audio_path: 音频文件路径
    :return: 时长（秒）
    """
    try:
        import ffmpeg
        probe = ffmpeg.probe(audio_path)
        duration = float(probe['format']['duration'])
        return duration
    except ImportError:
        # 如果没有 ffmpeg-python，使用 ffprobe 命令行
        try:
            import subprocess
            cmd = [
                "ffprobe", "-v", "quiet", "-print_format", "json",
                "-show_format", "-show_streams", audio_path
            ]
            result = subprocess.run(cmd, capture_output=True, text=True, check=True)
            probe = json.loads(result.stdout)
            duration = float(probe['format']['duration'])
            return duration
        except Exception as e:
            print(f"[WARN] 无法获取音频时长：{audio_path} ({e})")
            # 估算：中文普通话约 4-6 字/秒
            return 2.0  # 默认 2 秒
    except Exception as e:
        print(f"[WARN] 无法获取音频时长：{audio_path} ({e})")
        return 2.0


def calculate_timings(script: list, audio_dir: str = None, gap: float = 0.2) -> list:
    """
    计算时间轴
    
    :param script: 台词列表
    :param audio_dir: 配音文件目录（如果 audio_file 是相对路径）
    :param gap: 对话间隔（秒）
    :return: 带时间轴的台词列表
    """
    current_time = 0.0
    
    for item in script:
        # 获取配音文件时长
        audio_file = item.get("audio_file")
        
        if audio_file and os.path.isfile(audio_file):
            duration = get_audio_duration(audio_file)
        elif audio_file and audio_dir:
            audio_path = os.path.join(audio_dir, audio_file)
            if os.path.isfile(audio_path):
                duration = get_audio_duration(audio_path)
            else:
                # 估算时长
                text = item.get("text", "")
                duration = len(text) / 5.0  # 假设 5 字/秒
        else:
            # 没有配音文件，估算时长
            text = item.get("text", "")
            duration = len(text) / 5.0  # 假设 5 字/秒
        
        # 设置开始和结束时间
        item["start"] = current_time
        item["end"] = current_time + duration
        
        # 更新当前时间（加上间隔）
        current_time = item["end"] + gap
    
    print(f"[INFO] 时间轴计算完成：总时长 {current_time:.2f} 秒")
    return script


def generate_srt(script: list, output_path: str):
    """
    生成 SRT 字幕文件
    
    :param script: 带时间轴的台词列表
    :param output_path: 输出 SRT 文件路径
    """
    # 创建 pysrt SubRipFile
    subs = pysrt.SubRipFile()
    
    for item in script:
        # 创建字幕项
        sub = pysrt.SubRipItem()
        sub.index = item["id"]
        sub.text = item["text"]
        
        # 设置时间
        start_seconds = item["start"]
        end_seconds = item["end"]
        
        sub.start.hours = int(start_seconds // 3600)
        sub.start.minutes = int((start_seconds % 3600) // 60)
        sub.start.seconds = int(start_seconds % 60)
        sub.start.milliseconds = int((start_seconds % 1) * 1000)
        
        sub.end.hours = int(end_seconds // 3600)
        sub.end.minutes = int((end_seconds % 3600) // 60)
        sub.end.seconds = int(end_seconds % 60)
        sub.end.milliseconds = int((end_seconds % 1) * 1000)
        
        subs.append(sub)
    
    # 保存 SRT 文件
    os.makedirs(os.path.dirname(os.path.abspath(output_path)), exist_ok=True)
    subs.save(output_path, encoding="utf-8")
    
    print(f"[OK] SRT 字幕文件已生成：{output_path}")
    print(f"[INFO] 共 {len(script)} 条字幕")


def main():
    parser = argparse.ArgumentParser(
        description="Script to SRT with Timing · 剧本到 SRT 时间轴生成器",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
示例：
  python script-to-srt-with-timing.py --script script.json --output subtitles.srt
  python script-to-srt-with-timing.py --script script.md --audio-dir ./audio/ --output subtitles.srt
  python script-to-srt-with-timing.py --script script.json --gap 0.3 --output subtitles.srt
"""
    )
    parser.add_argument("--script", required=True, help="剧本文件路径（JSON 或 MD 格式）")
    parser.add_argument("--audio-dir", help="配音文件目录（如果 audio_file 是相对路径）")
    parser.add_argument("--output", required=True, help="输出 SRT 文件路径")
    parser.add_argument("--gap", type=float, default=0.2, help="对话间隔（秒，默认 0.2）")
    
    args = parser.parse_args()
    
    # 检查剧本文件是否存在
    if not os.path.isfile(args.script):
        print(f"[ERROR] 剧本文件不存在：{args.script}")
        sys.exit(1)
    
    # 解析剧本
    if args.script.lower().endswith(".json"):
        script = parse_script_json(args.script)
    elif args.script.lower().endswith(".md"):
        script = parse_script_md(args.script)
    else:
        print(f"[ERROR] 不支持的剧本格式：{args.script}")
        print("支持格式：.json, .md")
        sys.exit(1)
    
    if not script:
        print("[ERROR] 剧本为空，请检查格式")
        sys.exit(1)
    
    # 计算时间轴
    script = calculate_timings(script, args.audio_dir, args.gap)
    
    # 生成 SRT 文件
    generate_srt(script, args.output)
    
    print("\n[OK] 处理完成")


if __name__ == "__main__":
    main()