# ------------
# 音频生成与语速调整模块（简化版）
# 该模块负责TTS音频生成和简单的语速调整
# 主要功能包括：
# - 为每个翻译片段独立生成TTS音频
# - 计算所需的语速调整因子
# - 调整音频速度以匹配原始时长
# - 保持与翻译结果完全一致的分块结构
# ------------

import json
import os
import time
import shutil
import subprocess
from typing import Tuple

import pandas as pd
from pydub import AudioSegment
from rich.console import Console
from rich.progress import Progress

from core.utils import *
from core.utils.models import *
from core.asr_backend.audio_preprocess import get_audio_duration
from core.tts_backend.tts_main import tts_main

console = Console()

TEMP_FILE_TEMPLATE = f"{_AUDIO_TMP_DIR}/{{}}_temp.wav"
OUTPUT_FILE_TEMPLATE = f"{_AUDIO_SEGS_DIR}/{{}}.wav"

# ------------
# 音频速度调整函数
# 使用FFmpeg调整音频播放速度
# 参数: input_file - 输入音频文件路径
#       output_file - 输出音频文件路径  
#       speed_factor - 速度调整因子(>1加速，<1减速)
# ------------
def adjust_audio_speed(input_file: str, output_file: str, speed_factor: float) -> None:
    if abs(speed_factor - 1.0) < 0.001:
        shutil.copy2(input_file, output_file)
        return

    atempo_filters = []
    temp_speed = speed_factor
    while temp_speed > 2.0:
        atempo_filters.append("atempo=2.0")
        temp_speed /= 2.0
    atempo_filters.append(f"atempo={temp_speed}")
    
    filter_str = ",".join(atempo_filters)
    cmd = ['ffmpeg', '-i', input_file, '-filter:a', filter_str, '-y', output_file]

    input_duration = get_audio_duration(input_file)
    max_retries = 2
    
    for attempt in range(max_retries):
        try:
            subprocess.run(cmd, check=True, stderr=subprocess.PIPE, stdout=subprocess.DEVNULL)
            output_duration = get_audio_duration(output_file)
            expected_duration = input_duration / speed_factor
            diff = output_duration - expected_duration
            
            if output_duration >= expected_duration * 1.02 and diff <= 0.1:
                audio = AudioSegment.from_wav(output_file)
                trimmed_audio = audio[:(expected_duration * 1000)]
                trimmed_audio.export(output_file, format="wav")
                rprint(f"[yellow]✂️ Audio too long ({output_duration:.2f}s > {expected_duration:.2f}s), trimmed to expected duration.[/yellow]")
                return
            elif output_duration >= expected_duration * 1.02:
                raise Exception(f"Audio duration abnormal: input file={input_file}, output file={output_file}, speed factor={speed_factor}, input duration={input_duration:.2f}s, output duration={output_duration:.2f}s")
            return
        except subprocess.CalledProcessError as e:
            if attempt < max_retries - 1:
                rprint(f"[yellow]⚠️ Audio speed adjustment failed, retrying in 1s ({attempt + 1}/{max_retries})[/yellow]")
                time.sleep(1)
            else:
                rprint(f"[red]❌ Audio speed adjustment failed, max retries reached ({max_retries})[/red]")
                raise e

# ------------
# 简化的音频生成主控制函数
# 为每个翻译片段独立生成音频并进行语速调整
# 主要流程: 1. 创建必要的工作目录
#          2. 加载配音任务列表
#          3. 为每个片段生成TTS音频
#          4. 计算所需的语速调整因子
#          5. 调整音频速度以匹配原始时长
# ------------
def gen_audio() -> None:
    rprint("[bold magenta]🚀 Starting simple audio generation with speed adjustment...[/bold magenta]")
    
    # 创建工作目录
    os.makedirs(_AUDIO_TMP_DIR, exist_ok=True)
    os.makedirs(_AUDIO_SEGS_DIR, exist_ok=True)
    
    # 读取配音任务列表
    with open(_DUB_CHUNKS, 'r', encoding='utf-8') as f:
        dub_tasks = json.load(f)
    
    rprint(f"[green]📊 Loaded {len(dub_tasks)} dubbing tasks[/green]")
    
    # 处理每个片段
    with Progress() as progress:
        task = progress.add_task("[cyan]🔄 Processing audio segments...", total=len(dub_tasks))
        
        for dub_task in dub_tasks:
            task_id = dub_task['task_id']
            text = dub_task['text']
            original_duration = dub_task['duration']
            target_duration = dub_task['target_duration']
            
            rprint(f"[cyan]🎯 Processing segment {task_id}: '{text[:50]}...'[/cyan]")
            
            # 步骤1: 生成TTS音频
            temp_file = TEMP_FILE_TEMPLATE.format(task_id)
            temp_df = pd.DataFrame([{
                'number': task_id,
                'text': text
            }])
            
            try:
                tts_main(text, temp_file, task_id, temp_df)
                tts_duration = get_audio_duration(temp_file)
                
                rprint(f"[green]✅ TTS generated: {tts_duration:.2f}s (target: {target_duration:.2f}s)[/green]")
                
                # 步骤2: 计算语速调整因子
                speed_factor = tts_duration / target_duration
                print(f"cal_speed_factor: {speed_factor}")
                
                # 限制速度调整范围
                max_speed = 5.0  # 最大2倍速
                min_speed = 0.5  # 最小0.5倍速
                
                if speed_factor > max_speed:
                    rprint(f"[yellow]⚠️ Required speed {speed_factor:.2f} > max {max_speed}, capping at {max_speed}[/yellow]")
                    speed_factor = max_speed
                elif speed_factor < min_speed:
                    rprint(f"[yellow]⚠️ Required speed {speed_factor:.2f} < min {min_speed}, capping at {min_speed}[/yellow]")
                    speed_factor = min_speed
                
                print(f"speed_factor: {speed_factor}")

                # 步骤3: 调整音频速度
                output_file = OUTPUT_FILE_TEMPLATE.format(task_id)
                adjust_audio_speed(temp_file, output_file, speed_factor)
                
                final_duration = get_audio_duration(output_file)
                
                # 更新任务信息
                dub_task['real_dur'] = tts_duration
                dub_task['speed_factor'] = speed_factor
                dub_task['final_duration'] = final_duration
                dub_task['adjusted_start'] = dub_task['start_time']
                dub_task['adjusted_end'] = dub_task['start_time'] + final_duration

                print(dub_task)
                
                # 步骤4: 检查结果
                time_diff = final_duration - target_duration
                if abs(time_diff) <= 0.1:
                    emoji = "✅"
                    status = "Perfect"
                elif abs(time_diff) <= 0.3:
                    emoji = "⚡"
                    status = "Good"
                else:
                    emoji = "⚠️"
                    status = "Acceptable"
                
                rprint(f"[cyan]{emoji} Segment {task_id} {status}: {final_duration:.2f}s (diff: {time_diff:+.2f}s, speed: {speed_factor:.2f}x)[/cyan]")
                
                progress.advance(task)
                
            except Exception as e:
                rprint(f"[red]❌ Error processing segment {task_id}: {e}[/red]")
                raise e
    
    # 保存更新后的配音任务信息
    with open(_DUB_CHUNKS, 'w', encoding='utf-8') as f:
        json.dump(dub_tasks, f, ensure_ascii=False, indent=2)
    
    rprint("[bold green]🎉 Audio generation completed![/bold green]")
    
    # 显示统计信息
    total_segments = len(dub_tasks)
    perfect_matches = sum(1 for task in dub_tasks if abs(task.get('final_duration', 0) - task['target_duration']) <= 0.1)
    good_matches = sum(1 for task in dub_tasks if 0.1 < abs(task.get('final_duration', 0) - task['target_duration']) <= 0.3)
    
    rprint(f"[green]📊 Results: {perfect_matches}/{total_segments} perfect, {good_matches}/{total_segments} good[/green]")

if __name__ == "__main__":
    gen_audio()