# ------------
# 参考音频提取模块
# 该模块用于从原始音频中提取用于语音合成的参考音频片段
# 主要功能包括：
# - 时间格式转换为音频采样点
# - 音频片段提取和保存
# - 批量处理音频任务列表
# - 与Demucs音源分离模块集成
# ------------

import json
import os
from rich.panel import Panel
from rich.console import Console
from rich.progress import Progress, SpinnerColumn, TextColumn, BarColumn
from core.utils import *
from core.utils.models import *
import pandas as pd
import soundfile as sf
console = Console()
from core.asr_backend.demucs_vl import demucs_audio
from core.utils.models import *

# ------------
# 时间字符串转换为音频采样点数函数
# 用于将时间格式(HH:MM:SS,mmm)转换为对应的音频采样点位置
# 参数: time_str - 时间字符串格式如"00:01:30,500" 
#       sr - 音频采样率
# 返回: 对应的采样点数位置
# ------------
def time_to_samples(time_str, sr):
    h, m, s = time_str.split(':')
    s, ms = s.split(',') if ',' in s else (s, '0')
    seconds = int(h) * 3600 + int(m) * 60 + float(s) + float(ms) / 1000
    return int(seconds * sr)

# ------------
# 音频片段提取函数
# 从完整音频数据中提取指定时间段的音频片段并保存为文件
# 参数: audio_data - 完整音频数据数组
#       sr - 音频采样率
#       start_time - 开始时间字符串(HH:MM:SS,mmm格式)
#       end_time - 结束时间字符串(HH:MM:SS,mmm格式)  
#       out_file - 输出文件路径
# ------------
def extract_audio(audio_data, sr, start_time, end_time, out_file):
    # 重构后：支持直接的秒数格式
    if isinstance(start_time, (int, float)):
        start = int(start_time * sr)
        end = int(end_time * sr)
    else:
        # 保持对原有字符串格式的兼容
        start = time_to_samples(start_time, sr)
        end = time_to_samples(end_time, sr)
    sf.write(out_file, audio_data[start:end], sr)

# ------------
# 参考音频提取主函数
# 该函数负责从原始音频中提取用于语音合成的参考音频片段
# 主要流程：
# 1. 运行Demucs音源分离(如果尚未运行)
# 2. 检查是否已存在音频片段，避免重复处理
# 3. 创建输出目录存放参考音频
# 4. 读取音频任务配置文件获取时间段信息
# 5. 加载分离后的人声音频文件
# 6. 循环处理每个时间段，提取对应的音频片段
# 7. 显示处理进度并保存结果
# ------------
def extract_refer_audio_main():
    demucs_audio()
    if os.path.exists(os.path.join(_AUDIO_SEGS_DIR, '1.wav')):
        rprint(Panel("Audio segments already exist, skipping extraction", title="Info", border_style="blue"))
        return

    os.makedirs(_AUDIO_REFERS_DIR, exist_ok=True)
    
    # 重构后使用JSON格式的音频任务
    with open(_AUDIO_TASKS, 'r', encoding='utf-8') as f:
        audio_tasks = json.load(f)
    
    # 同时读取翻译结果以获取时间戳信息
    with open(_TRANSLATE_RESULT, 'r', encoding='utf-8') as f:
        translate_result = json.load(f)
    
    data, sr = sf.read(_VOCAL_AUDIO_FILE)
    
    with Progress(
        SpinnerColumn(),
        TextColumn("[progress.description]{task.description}"),
        BarColumn(),
        TextColumn("[progress.percentage]{task.percentage:>3.0f}%"),
    ) as progress:
        task = progress.add_task("Extracting audio segments...", total=len(translate_result))
        
        for segment in translate_result:
            segment_id = segment['id']
            start_time = segment['start']
            end_time = segment['end']
            out_file = os.path.join(_AUDIO_REFERS_DIR, f"{segment_id}.wav")
            extract_audio(data, sr, start_time, end_time, out_file)
            progress.update(task, advance=1)
            
    rprint(Panel(f"Audio segments saved to {_AUDIO_REFERS_DIR}", title="Success", border_style="green"))

if __name__ == "__main__":
    extract_refer_audio_main()