#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
根据字幕的startpos和endpos截取中文音频
用法: python extract_cn_audio.py <adid>
"""

import os
import sys
import sqlite3
import subprocess
import shutil


def get_file_duration(file_path):
    """
    获取音频文件时长（毫秒）
    """
    cmd = ["ffprobe", "-v", "quiet", "-show_entries", "format=duration", "-of", "default=noprint_wrappers=1:nokey=1", file_path]
    result = subprocess.run(cmd, capture_output=True, text=True)
    if result.returncode == 0:
        return int(float(result.stdout.strip()) * 1000)
    return 0


def get_subtitles(adid, db_path):
    """
    从数据库获取字幕信息
    """
    conn = sqlite3.connect(db_path)
    cursor = conn.cursor()
    
    cursor.execute("""
        SELECT idx, start, end, cn FROM srts 
        WHERE adid = ? 
        ORDER BY idx
    """, (adid,))
    
    subtitles = []
    for row in cursor.fetchall():
        idx, start, end, cn = row
        # 解析时间格式
        start_ms = int(start)
        end_ms = int(end)
        subtitles.append((idx, start_ms, end_ms, cn))
    
    conn.close()
    return subtitles


def get_adname(adid, db_path):
    """
    从数据库获取adname
    """
    conn = sqlite3.connect(db_path)
    cursor = conn.cursor()
    
    cursor.execute("SELECT DISTINCT adname FROM srts WHERE adid = ? LIMIT 1", (adid,))
    result = cursor.fetchone()
    conn.close()
    
    if result:
        return result[0]
    return f"adid_{adid}"


def extract_cn_audio(adid):
    """
    根据startpos和endpos截取中文音频
    """
    # 配置
    db_path = r'\\ll\D\xtrssvjj\db\dpd.db'
    
    # 获取adname
    adname = get_adname(adid, db_path)
    print(f"adname: {adname}")
    
    # 音频文件路径
    audio_file = os.path.join(r'\\ll\D\xtrssvjj', f"{adname}.mp3")
    print(f"音频文件: {audio_file}")
    
    # 检查音频文件是否存在
    if not os.path.exists(audio_file):
        print(f"错误：音频文件不存在: {audio_file}")
        return False
    
    # 获取音频文件总长度
    audio_duration = get_file_duration(audio_file)
    print(f"音频总长度: {audio_duration} 毫秒")
    
    # 输出目录
    output_dir = f"{adname}_cn_audio"
    
    # 清理输出目录
    if os.path.exists(output_dir):
        shutil.rmtree(output_dir)
    os.makedirs(output_dir)
    
    # 获取字幕
    subtitles = get_subtitles(adid, db_path)
    print(f"找到 {len(subtitles)} 句字幕")
    
    if not subtitles:
        print("错误：未找到字幕信息")
        return False
    
    # 计算每句的startpos和endpos
    enhanced_subtitles = []
    for i, (idx, start, end, cn) in enumerate(subtitles):
        prevend = subtitles[i-1][2] if i > 0 else 0
        startgap = start - prevend if i > 0 else 0
        
        if i < len(subtitles)-1:
            nextstart = subtitles[i+1][1]
            endgap = nextstart - end
        else:
            # 最后一句的nextstart为音频文件的实际总长度
            nextstart = audio_duration
            endgap = nextstart - end if audio_duration > end else 0
        
        # 计算startpos和endpos
        # 第一句的startpos固定为0
        if i == 0:
            startpos = 0
        else:
            startpos = (prevend + start) // 2
        
        # 最后一句的endpos为视频的最后位置(总长度-1)
        if i == len(subtitles)-1:
            endpos = audio_duration - 1
        else:
            endpos = (end + nextstart) // 2
        
        # 计算addstart值：第一句addstart=0，后面的用400-(start-startpos)
        if i == 0:
            addstart = 0
        else:
            addstart = 400 - (start - startpos)
            if addstart <= 0:
                addstart = ''
        
        # 计算addend值：最后一句addend=0，其他用400-(endpos-end)
        if i == len(subtitles)-1:
            addend = 0
        else:
            addend = 400 - (endpos - end)
            if addend <= 0:
                addend = ''
        
        enhanced_subtitles.append((idx, start, end, startpos, endpos, addstart, addend, cn))
    
    # 截取音频
    for i, (idx, start, end, startpos, endpos, addstart, addend, cn) in enumerate(enhanced_subtitles):
        # 转换为秒
        start_sec = startpos / 1000
        duration_sec = (endpos - startpos) / 1000
        
        # 输出文件名
        output_path = os.path.join(output_dir, f"seg{idx:03d}.mp3")
        
        # 执行截取
        cmd = [
            "ffmpeg", "-y",
            "-ss", str(start_sec),
            "-t", str(duration_sec),
            "-i", audio_file,
            "-acodec", "libmp3lame",
            "-ar", "44100",
            "-ac", "2",
            output_path
        ]
        
        subprocess.run(cmd, capture_output=True)
        
        # 获取实际音频长度
        actual_duration = get_file_duration(output_path)
        
        # 重命名文件，包含更多信息
        new_filename = f"seg{idx:03d}-{startpos}-{endpos}-{actual_duration}.mp3"
        new_path = os.path.join(output_dir, new_filename)
        
        # 处理文件重命名冲突
        if os.path.exists(new_path):
            os.remove(new_path)
        os.rename(output_path, new_path)
        
        # 显示进度
        if (i + 1) % 10 == 0:
            print(f"截取音频 {i+1}/{len(enhanced_subtitles)}")
    
    print(f"完成！截取了 {len(enhanced_subtitles)} 个音频片段")
    print(f"输出目录: {output_dir}")
    return True


def main():
    """
    主函数
    """
    if len(sys.argv) != 2:
        print("用法: python extract_cn_audio.py <adid>")
        print("示例: python extract_cn_audio.py 8496")
        sys.exit(1)
    
    try:
        adid = int(sys.argv[1])
    except ValueError:
        print("错误：adid必须是数字")
        sys.exit(1)
    
    print(f"=== 中文音频截取工具 ===")
    print(f"adid: {adid}")
    
    success = extract_cn_audio(adid)
    
    if success:
        print("\n截取完成！")
    else:
        print("\n截取失败！")
        sys.exit(1)


if __name__ == "__main__":
    main()