import pandas as pd
import json
import concurrent.futures
from core.translate_lines import translate_lines
from core._4_1_summarize import search_things_to_note_in_prompt
# 重构后不再需要这些导入
# from core._8_1_audio_task import check_len_then_trim
# from core._6_gen_sub import align_timestamp
from core.utils import *
from rich.console import Console
from rich.progress import Progress, SpinnerColumn, TextColumn
from difflib import SequenceMatcher
from core.utils.models import *
console = Console()

# 🔍 Translate a single ASR segment
def translate_segment(segment, theme_prompt, segments_list, segment_index):
    """
    翻译单个语义切割片段
    
    Args:
        segment: 语义切割片段，包含id, text, start, end
        theme_prompt: 主题提示
        segments_list: 完整的片段列表，用于提供上下文
        segment_index: 当前片段在列表中的索引
    """
    text = segment['text']
    things_to_note_prompt = search_things_to_note_in_prompt(text)
    
    # 提供前后文本作为上下文
    previous_content_prompt = None
    after_content_prompt = None
    
    if segment_index > 0:
        # 获取前一个片段的文本
        previous_content_prompt = [segments_list[segment_index - 1]['text']]
    
    if segment_index < len(segments_list) - 1:
        # 获取后一个片段的文本
        after_content_prompt = [segments_list[segment_index + 1]['text']]
    
    translation, english_result = translate_lines(text, previous_content_prompt, after_content_prompt, things_to_note_prompt, theme_prompt, segment_index)
    return segment['id'], text, translation

# Add similarity calculation function
def similar(a, b):
    return SequenceMatcher(None, a, b).ratio()

# 🚀 Main function to translate all chunks
@check_file_exists(_TRANSLATE_RESULT)
def translate_all():
    console.print("[bold green]Start Translating All...[/bold green]")
    
    # 读取语义切割结果而不是ASR原始结果
    with open(_SPLIT_RESULT, 'r', encoding='utf-8') as f:
        asr_segments = json.load(f)
    
    with open(_4_1_TERMINOLOGY, 'r', encoding='utf-8') as file:
        theme_prompt = json.load(file).get('theme')

    # 🔄 Use concurrent execution for translation
    with Progress(SpinnerColumn(), TextColumn("[progress.description]{task.description}"), transient=True) as progress:
        task = progress.add_task("[cyan]Translating segments...", total=len(asr_segments))
        with concurrent.futures.ThreadPoolExecutor(max_workers=load_key("max_workers")) as executor:
            futures = []
            for i, segment in enumerate(asr_segments):
                future = executor.submit(translate_segment, segment, theme_prompt, asr_segments, i)
                futures.append(future)
            results = []
            for future in concurrent.futures.as_completed(futures):
                results.append(future.result())
                progress.update(task, advance=1)

    results.sort(key=lambda x: x[0])  # Sort results based on original order
    
    # 💾 创建翻译结果数据结构
    translate_result = []
    
    # 按顺序处理翻译结果，并与原语义切割片段合并
    for segment_id, source_text, translated_text in results:
        # 在asr_segments中找到对应的片段，获取时间戳
        original_segment = next(seg for seg in asr_segments if seg['id'] == segment_id)
        
        translate_result.append({
            "id": segment_id,
            "source_text": source_text,
            "translated_text": translated_text,
            "start": original_segment['start'],
            "end": original_segment['end']
        })
    
    # 保存翻译结果到JSON文件
    with open(_TRANSLATE_RESULT, 'w', encoding='utf-8') as f:
        json.dump(translate_result, f, ensure_ascii=False, indent=2)
    
    console.print(f"[bold green]✅ Translation completed and results saved to {_TRANSLATE_RESULT}.[/bold green]")
    console.print(f"[bold green]📊 Translated {len(translate_result)} segments.[/bold green]")

if __name__ == '__main__':
    translate_all()