#!/usr/bin/env python3
"""
测量所有TTS音频文件的时长，并保存到中间文件
"""
import os
import json
import subprocess

# 导入编码处理模块
from console_helper import (
    safe_print, print_header, print_step, print_info,
    print_success, print_fail, print_done
)

def get_audio_duration(audio_path):
    """获取音频时长"""
    args = [
        'ffprobe',
        '-v', 'quiet',
        '-print_format', 'json',
        '-show_format',
        audio_path
    ]
    
    try:
        result = subprocess.run(args, check=True, capture_output=True, text=True)
        import json
        data = json.loads(result.stdout)
        return float(data['format']['duration'])
    except Exception as e:
        safe_print(f"      获取音频时长失败: {e}")
        return 0

def main():
    audio_dir = "tts_audio_all"
    output_file = "audio_durations.json"
    
    print_header("测量音频时长")
    safe_print(f"音频目录: {audio_dir}")
    safe_print("")
    
    # 检查目录是否存在
    if not os.path.exists(audio_dir):
        safe_print(f"错误: 目录 {audio_dir} 不存在")
        return
    
    # 获取所有音频文件
    audio_files = sorted([f for f in os.listdir(audio_dir) if f.endswith('.wav')])
    
    if not audio_files:
        safe_print(f"错误: 目录 {audio_dir} 中没有找到音频文件")
        return
    
    print_info("找到音频文件", str(len(audio_files)))
    safe_print("")
    
    # 测量每个音频文件的时长
    print_step(1, 2, "测量音频时长...")
    durations = {}
    total_duration = 0
    
    for i, audio_file in enumerate(audio_files, 1):
        audio_path = os.path.join(audio_dir, audio_file)
        duration = get_audio_duration(audio_path)
        
        # 提取行号
        line_number = audio_file.replace('tts_line_', '').replace('.wav', '')
        durations[line_number] = {
            'file': audio_file,
            'duration': round(duration, 3),
            'duration_formatted': f"{int(duration // 60):02d}:{int(duration % 60):02d}.{int((duration % 1) * 1000):03d}"
        }
        total_duration += duration
        
        print_info(f"第{i}/{len(audio_files)}个", f"{audio_file}: {duration:.3f}秒")
    
    safe_print("")
    print_info("总时长", f"{total_duration:.2f}秒 ({total_duration/60:.2f}分钟)")
    safe_print("")
    
    # 保存到JSON文件
    print_step(2, 2, "保存到中间文件...")
    output_data = {
        'total_files': len(audio_files),
        'total_duration': round(total_duration, 3),
        'total_duration_formatted': f"{int(total_duration // 60):02d}:{int(total_duration % 60):02d}.{int((total_duration % 1) * 1000):03d}",
        'durations': durations
    }
    
    try:
        with open(output_file, 'w', encoding='utf-8') as f:
            json.dump(output_data, f, ensure_ascii=False, indent=2)
        print_info("已保存", output_file)
    except Exception as e:
        safe_print(f"保存文件失败: {e}")
        return
    
    print_done()

if __name__ == '__main__':
    main()