# LIB · TTS 配音模块 + 自动字幕生成
# D180 · 2026-07-10
"""Edge-TTS 配音引擎 · 多角色音色 · 自动SRT字幕"""
import asyncio, edge_tts, os, json, shutil

# 角色音色表 · 从 engines/tts-engine.py 加载（单一真相源）
try:
    import sys, os as _os
    _engine_dir = _os.path.join(_os.path.dirname(_os.path.dirname(_os.path.abspath(__file__))), "engines")
    if _engine_dir not in sys.path: sys.path.insert(0, _engine_dir)
    from tts_engine import DEFAULT_VOICES as _ENGINE_VOICES
    VOICES = {name: cfg["voice"] for name, cfg in _ENGINE_VOICES.items()}
except Exception:
    VOICES = {
        "苏白": "zh-CN-XiaoxiaoNeural",
        "诸葛风": "zh-CN-YunxiNeural",
        "萧灵汐": "zh-CN-XiaoyiNeural",
        "王执事": "zh-CN-YunyangNeural",
        "林昊": "zh-CN-YunxiNeural",
        "旁白": "zh-CN-YunyangNeural",
        "默认女": "zh-CN-XiaoxiaoNeural",
        "默认男": "zh-CN-YunxiNeural",
    }

FFMPEG = shutil.which("ffmpeg") or os.path.expanduser("~/.workbuddy/binaries/ffmpeg.exe")

def tts(text, output_path, voice="zh-CN-YunxiNeural", rate="+0%"):
    """生成语音 mp3（返回时间戳用于字幕）"""
    os.makedirs(os.path.dirname(output_path) or ".", exist_ok=True)
    words = []
    async def _run():
        comm = edge_tts.Communicate(text, voice, rate=rate)
        sub = edge_tts.SubMaker()
        with open(output_path, "wb") as f:
            async for chunk in comm.stream():
                if chunk["type"] == "audio":
                    f.write(chunk["data"])
                elif chunk["type"] == "SentenceBoundary":
                    words.append({
                        "text": chunk["text"],
                        "start": chunk["offset"] / 10000000,
                        "end": (chunk["offset"] + chunk["duration"]) / 10000000
                    })
    asyncio.run(_run())
    return {"path": output_path, "words": words} if os.path.exists(output_path) else None

def tts_for_character(text, character, output_path, rate="+0%"):
    """按角色名查音色生成"""
    voice = VOICES.get(character, "zh-CN-YunxiNeural")
    return tts(text, output_path, voice, rate)

def words_to_srt(words, output_path, offset=0):
    """句时间戳 → SRT 字幕文件"""
    if not words:
        return None
    lines = []
    for i, w in enumerate(words):
        def fmt(t):
            ms = int((offset + t) * 1000)
            h = ms // 3600000; ms %= 3600000
            m = ms // 60000; ms %= 60000
            s = ms // 1000; ms %= 1000
            return f"{h:02d}:{m:02d}:{s:02d},{ms:03d}"
        lines.append(f"{i+1}")
        lines.append(f"{fmt(w['start'])} --> {fmt(w['end'])}")
        lines.append(w["text"])
        lines.append("")
    os.makedirs(os.path.dirname(output_path) or ".", exist_ok=True)
    with open(output_path, "w", encoding="utf-8") as f:
        f.write("\n".join(lines))
    return output_path

def burn_subtitle(video_path, srt_path, output_path):
    """烧录字幕到视频"""
    if not os.path.exists(FFMPEG):
        return f"ffmpeg not found at {FFMPEG}"
    os.makedirs(os.path.dirname(output_path) or ".", exist_ok=True)
    import subprocess
    srt_esc = srt_path.replace("\\", "/").replace(":", "\\:")
    r = subprocess.run([
        FFMPEG, "-y",
        "-i", video_path,
        "-vf", f"subtitles='{srt_esc}':force_style='Fontsize=24,Alignment=2,BorderStyle=3,Outline=1,Shadow=1'",
        "-c:a", "copy",
        output_path
    ], capture_output=True, text=True)
    return output_path if r.returncode == 0 else r.stderr[:300]

def mix_audio_video(video_path, audio_path, output_path, video_vol=0.3):
    """FFmpeg 合成视频+音频"""
    if not os.path.exists(FFMPEG):
        return f"ffmpeg not found at {FFMPEG}"
    os.makedirs(os.path.dirname(output_path) or ".", exist_ok=True)
    import subprocess
    r = subprocess.run([
        FFMPEG, "-y",
        "-i", video_path,
        "-i", audio_path,
        "-filter_complex", f"[0:a]volume={video_vol}[va];[1:a]volume=1.0[aa];[va][aa]amix=inputs=2:duration=first",
        "-c:v", "copy",
        output_path
    ], capture_output=True, text=True)
    return output_path if r.returncode == 0 else r.stderr

def add_subtitle(video_path, subtitle_text, output_path):
    """FFmpeg 添加字幕"""
    if not os.path.exists(FFMPEG):
        return f"ffmpeg not found at {FFMPEG}"
    import subprocess
    r = subprocess.run([
        FFMPEG, "-y",
        "-i", video_path,
        "-vf", f"drawtext=text='{subtitle_text}':fontsize=28:fontcolor=white:borderw=2:bordercolor=black:x=(w-text_w)/2:y=h-120",
        "-c:a", "copy",
        output_path
    ], capture_output=True, text=True)
    return output_path if r.returncode == 0 else r.stderr

if __name__ == "__main__":
    import sys
    if len(sys.argv) < 2:
        print("Usage: python tts.py tts <text> <output> [voice]")
        print("       python tts.py char <text> <character> <output>")
        print("       python tts.py mix <video> <audio> <output>")
        print(f"\nVoices: {list(VOICES.keys())}")
        sys.exit(1)
    cmd = sys.argv[1]
    if cmd == "tts":
        voice = sys.argv[4] if len(sys.argv) > 4 else "zh-CN-YunxiNeural"
        print(tts(sys.argv[2], sys.argv[3], voice) or "FAILED")
    elif cmd == "char":
        print(tts_for_character(sys.argv[2], sys.argv[3], sys.argv[4]) or "FAILED")
    elif cmd == "mix":
        print(mix_audio_video(sys.argv[2], sys.argv[3], sys.argv[4]) or "FAILED")