import json
import os
import pandas as pd
import subprocess
from pydub import AudioSegment
from rich.progress import Progress, SpinnerColumn, TextColumn, BarColumn, TaskProgressColumn
from rich.console import Console
from core.utils import *
from core.utils.models import *
console = Console()

DUB_VOCAL_FILE = 'output/dub.mp3'

DUB_SUB_FILE = 'output/dub.srt'
OUTPUT_FILE_TEMPLATE = f"{_AUDIO_SEGS_DIR}/{{}}.wav"

def load_and_flatten_data(excel_file):
    """Load and flatten Excel data"""
    df = pd.read_excel(excel_file)
    lines = [eval(line) if isinstance(line, str) else line for line in df['lines'].tolist()]
    lines = [item for sublist in lines for item in sublist]
    
    new_sub_times = [eval(time) if isinstance(time, str) else time for time in df['new_sub_times'].tolist()]
    new_sub_times = [item for sublist in new_sub_times for item in sublist]
    
    return df, lines, new_sub_times

def get_audio_files(df):
    """Generate a list of audio file paths"""
    audios = []
    for index, row in df.iterrows():
        number = row['number']
        line_count = len(eval(row['lines']) if isinstance(row['lines'], str) else row['lines'])
        for line_index in range(line_count):
            temp_file = OUTPUT_FILE_TEMPLATE.format(f"{number}_{line_index}")
            audios.append(temp_file)
    return audios

def process_audio_segment(audio_file):
    """Process a single audio segment with MP3 compression"""
    temp_file = f"{audio_file}_temp.mp3"
    ffmpeg_cmd = [
        'ffmpeg', '-y',
        '-i', audio_file,
        '-ar', '16000',
        '-ac', '1',
        '-b:a', '64k',
        temp_file
    ]
    subprocess.run(ffmpeg_cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
    audio_segment = AudioSegment.from_mp3(temp_file)
    os.remove(temp_file)
    return audio_segment

def merge_audio_segments(audios, new_sub_times, sample_rate):
    merged_audio = AudioSegment.silent(duration=0, frame_rate=sample_rate)
    
    with Progress(SpinnerColumn(), TextColumn("[progress.description]{task.description}"), BarColumn(), TaskProgressColumn()) as progress:
        merge_task = progress.add_task("🎵 Merging audio segments...", total=len(audios))
        
        for i, (audio_file, time_range) in enumerate(zip(audios, new_sub_times)):
            if not os.path.exists(audio_file):
                console.print(f"[bold yellow]⚠️  Warning: File {audio_file} does not exist, skipping...[/bold yellow]")
                progress.advance(merge_task)
                continue
                
            audio_segment = process_audio_segment(audio_file)
            start_time, end_time = time_range
            
            # Add silence segment
            if i > 0:
                prev_end = new_sub_times[i-1][1]
                silence_duration = start_time - prev_end
                if silence_duration > 0:
                    silence = AudioSegment.silent(duration=int(silence_duration * 1000), frame_rate=sample_rate)
                    merged_audio += silence
            elif start_time > 0:
                silence = AudioSegment.silent(duration=int(start_time * 1000), frame_rate=sample_rate)
                merged_audio += silence
                
            merged_audio += audio_segment
            progress.advance(merge_task)
    
    return merged_audio

def create_srt_subtitle():
    """基于配音块数据创建字幕文件"""
    # 读取配音块数据
    with open(_DUB_CHUNKS, 'r', encoding='utf-8') as f:
        dub_chunks = json.load(f)
    
    with open(DUB_SUB_FILE, 'w', encoding='utf-8') as f:
        for i, chunk in enumerate(dub_chunks, 1):
            # 使用调整后的时间戳（如果存在），否则使用原始时间戳
            start_time = chunk.get('adjusted_start', chunk['start_time'])
            end_time = chunk.get('adjusted_end', chunk['end_time'])
            text = chunk['text']
            
            start_str = f"{int(start_time//3600):02d}:{int((start_time%3600)//60):02d}:{int(start_time%60):02d},{int((start_time*1000)%1000):03d}"
            end_str = f"{int(end_time//3600):02d}:{int((end_time%3600)//60):02d}:{int(end_time%60):02d},{int((end_time*1000)%1000):03d}"
            
            f.write(f"{i}\n")
            f.write(f"{start_str} --> {end_str}\n")
            f.write(f"{text}\n\n")
    
    rprint(f"[bold green]✅ Subtitle file created: {DUB_SUB_FILE}[/bold green]")

def merge_full_audio():
    """Main function: Process the complete audio merging process"""
    console.print("\n[bold cyan]🎬 Starting audio merging process...[/bold cyan]")
    
    # 读取配音块数据以获取调整后的时间戳信息
    with open(_DUB_CHUNKS, 'r', encoding='utf-8') as f:
        dub_chunks = json.load(f)
    
    console.print(f"[bold green]✅ Loaded {len(dub_chunks)} dubbing segments[/bold green]")
    
    # 获取视频总时长（基于最后一个片段的结束时间）
    total_duration = max(chunk['end_time'] for chunk in dub_chunks)
    console.print(f"[bold green]📏 Total video duration: {total_duration:.2f}s[/bold green]")
    
    # 创建静音的完整音轨
    sample_rate = 16000
    merged_audio = AudioSegment.silent(duration=int(total_duration * 1000), frame_rate=sample_rate)
    
    console.print("[bold cyan]🔄 Starting audio merge process...[/bold cyan]")
    
    # 按照调整后的时间戳将音频片段合并到完整音轨
    with Progress(SpinnerColumn(), TextColumn("[progress.description]{task.description}"), BarColumn(), TaskProgressColumn()) as progress:
        merge_task = progress.add_task("🎵 Merging audio segments...", total=len(dub_chunks))
        
        for chunk in dub_chunks:
            task_id = chunk['task_id']
            
            # 使用调整后的时间戳（如果存在），否则使用原始时间戳
            start_time = chunk.get('adjusted_start', chunk['start_time'])
            
            # 找到对应的音频文件
            audio_filename = f"{task_id}.wav"
            audio_path = os.path.join(_AUDIO_SEGS_DIR, audio_filename)
            
            if os.path.exists(audio_path):
                try:
                    # 加载音频片段，自动检测格式
                    audio_segment = AudioSegment.from_file(audio_path)
                    # 将音频片段叠加到指定时间位置
                    start_time_ms = int(start_time * 1000)
                    merged_audio = merged_audio.overlay(audio_segment, position=start_time_ms)
                    progress.advance(merge_task)
                except Exception as e:
                    console.print(f"[bold yellow]⚠️  Warning: Failed to load audio file {audio_path}: {e}[/bold yellow]")
                    progress.advance(merge_task)
            else:
                console.print(f"[bold yellow]⚠️  Warning: Audio file {audio_path} not found, skipping...[/bold yellow]")
                progress.advance(merge_task)
    
    # 导出最终音频文件
    with console.status("[bold cyan]💾 Exporting final audio file...[/bold cyan]"):
        merged_audio = merged_audio.set_frame_rate(sample_rate).set_channels(1)
        merged_audio.export(_DUB_FINAL, format="wav")
    
    console.print(f"[bold green]✅ Audio file successfully merged![/bold green]")
    console.print(f"[bold green]📁 Output file: {_DUB_FINAL}[/bold green]")

if __name__ == "__main__":
    merge_full_audio()