# ------------
# 导入所需模块和库
# ------------
import datetime                                                                  # 导入日期时间处理模块
import json                                                                      # 导入JSON处理模块
import os                                                                        # 导入操作系统模块
import re                                                                        # 导入正则表达式模块
import pandas as pd                                                              # 导入数据分析库
from rich.console import Console                                                 # 导入富文本控制台输出库
from rich.panel import Panel                                                     # 导入富文本面板组件
from core.prompts import get_subtitle_trim_prompt                               # 导入字幕修剪提示生成函数
from core.tts_backend.estimate_duration import init_estimator, estimate_duration  # 导入时长估算相关函数
from core.utils import *                                                         # 导入核心工具函数
from core.utils.models import *                                                  # 导入模型相关工具

# ------------
# 全局变量和常量定义
# ------------
console = Console()                                                              # 初始化控制台对象
speed_factor = load_key("speed_factor")                                         # 加载语速倍数配置

TRANS_SUBS_FOR_AUDIO_FILE = 'output/audio/trans_subs_for_audio.srt'             # 翻译字幕文件路径
SRC_SUBS_FOR_AUDIO_FILE = 'output/audio/src_subs_for_audio.srt'                 # 原始字幕文件路径
ESTIMATOR = None                                                                 # 时长估算器全局变量

# # ------------
# # 字幕长度检查和修剪函数
# # ------------
# def check_len_then_trim(text, duration):
#     """
#     检查字幕文本的朗读时长，如果超过给定时长则进行修剪。
    
#     该函数首先使用时长估算器评估文本的朗读时间，如果估算时长超过给定时长，
#     则通过AI模型对文本进行智能修剪，如果AI调用失败则使用简单的标点符号移除方法。
    
#     Args:
#         text (str): 需要检查的字幕文本内容
#         duration (float): 字幕的最大允许时长（秒）
    
#     Returns:
#         str: 修剪后的字幕文本，如果原文本时长合适则返回原文本
    
#     Global Variables:
#         ESTIMATOR: 全局时长估算器实例，用于预测文本朗读时间
    
#     Notes:
#         - 估算时长会除以最大语速倍数以获得更准确的预测
#         - 如果AI拒绝修剪（敏感内容），会自动降级到简单的标点符号移除
#         - 修剪过程会在控制台显示详细的前后对比信息
#     """
#     global ESTIMATOR                                                             # 使用全局估算器变量
#     if ESTIMATOR is None:                                                        # 如果估算器未初始化
#         ESTIMATOR = init_estimator()                                             # 初始化时长估算器
#     estimated_duration = estimate_duration(text, ESTIMATOR) / speed_factor['max']  # 估算文本朗读时长，除以最大语速倍数
    
#     console.print(f"Subtitle text: {text}, "                                    # 打印字幕文本信息
#                   f"[bold green]Estimated reading duration: {estimated_duration:.2f} seconds[/bold green]")  # 打印估算时长

#     if estimated_duration > duration:                                            # 如果估算时长超过给定时长
#         rprint(Panel(f"Estimated reading duration {estimated_duration:.2f} seconds exceeds given duration {duration:.2f} seconds, shortening...", title="Processing", border_style="yellow"))  # 显示需要缩短的警告
#         original_text = text                                                     # 保存原始文本
#         prompt = get_subtitle_trim_prompt(text, duration)                        # 获取字幕修剪的AI提示
#         def valid_trim(response):                                                # 定义响应验证函数
#             if 'result' not in response:                                         # 如果响应中没有result字段
#                 return {'status': 'error', 'message': 'No result in response'}   # 返回错误状态
#             return {'status': 'success', 'message': ''}                          # 返回成功状态
#         try:                                                                     # 尝试使用AI进行修剪
#             response = ask_gpt(prompt, resp_type='json', log_title='sub_trim', valid_def=valid_trim)  # 调用GPT进行字幕修剪
#             shortened_text = response['result']                                  # 获取修剪后的文本
#         except Exception:                                                        # 如果AI调用失败
#             rprint("[bold red]🚫 AI refused to answer due to sensitivity, so manually remove punctuation[/bold red]")  # 打印AI拒绝信息
#             shortened_text = re.sub(r'[,.!?;:，。！？；：]', ' ', text).strip()  # 手动移除标点符号进行简单修剪
#         rprint(Panel(f"Subtitle before shortening: {original_text}\nSubtitle after shortening: {shortened_text}", title="Subtitle Shortening Result", border_style="green"))  # 显示修剪前后对比
#         return shortened_text                                                    # 返回修剪后的文本
#     else:                                                                        # 如果时长合适
#         return text                                                              # 直接返回原文本

# ------------
# 时间差计算函数
# ------------
def time_diff_seconds(t1, t2, base_date):
    """
    计算两个时间对象之间的秒数差值。
    
    由于time对象不包含日期信息，需要与基准日期组合成完整的datetime对象
    才能进行时间差计算。主要用于计算字幕的持续时间或字幕间的时间间隔。
    
    Args:
        t1 (datetime.time): 起始时间对象
        t2 (datetime.time): 结束时间对象  
        base_date (datetime.date): 基准日期，用于组合成完整的datetime对象
    
    Returns:
        float: 两个时间之间的秒数差值，可以是小数（包含毫秒）
    
    Example:
        >>> from datetime import time, date
        >>> t1 = time(10, 30, 15, 500000)  # 10:30:15.500
        >>> t2 = time(10, 30, 18, 200000)  # 10:30:18.200  
        >>> base = date.today()
        >>> time_diff_seconds(t1, t2, base)
        2.7
    """
    dt1 = datetime.datetime.combine(base_date, t1)                               # 将第一个时间对象与基准日期组合成完整的日期时间
    dt2 = datetime.datetime.combine(base_date, t2)                               # 将第二个时间对象与基准日期组合成完整的日期时间
    return (dt2 - dt1).total_seconds()                                           # 计算两个日期时间之间的总秒数差

# # ------------
# # SRT字幕文件处理函数
# # ------------
# def process_srt():
#     """
#     处理SRT字幕文件，生成用于音频合成的结构化数据。
    
#     该函数完成以下主要任务：
#     1. 读取翻译字幕文件和原始字幕文件
#     2. 解析SRT格式的字幕块（序号、时间戳、文本内容）
#     3. 清理文本内容（移除括号、特殊字符等）
#     4. 合并时长过短的相邻字幕块
#     5. 延长过短字幕的显示时间至最小时长
#     6. 格式化时间戳为标准字符串格式
    
#     文件路径:
#         - 翻译字幕: output/audio/trans_subs_for_audio.srt
#         - 原始字幕: output/audio/src_subs_for_audio.srt
    
#     Returns:
#         pandas.DataFrame: 包含以下列的数据框：
#             - number (int): 字幕序号
#             - start_time (str): 开始时间戳 (HH:MM:SS.mmm格式)
#             - end_time (str): 结束时间戳 (HH:MM:SS.mmm格式) 
#             - duration (float): 持续时间（秒）
#             - text (str): 清理后的翻译文本
#             - origin (str): 对应的原始文本
    
#     Raises:
#         FileNotFoundError: 当字幕文件不存在时
#         ValueError: 当字幕格式解析失败时（会跳过错误块继续处理）
    
#     Notes:
#         - 自动处理最小字幕时长限制，确保TTS有足够时间朗读
#         - 清理文本时会移除可能导致TTS错误的特殊字符
#         - 合并策略优先考虑相邻字幕的时间间隔
#     """
    
#     # ------------
#     # 读取翻译和原始字幕文件
#     # ------------
#     with open(TRANS_SUBS_FOR_AUDIO_FILE, 'r', encoding='utf-8') as file:         # 打开翻译字幕文件
#         content = file.read()                                                    # 读取翻译字幕内容
    
#     with open(SRC_SUBS_FOR_AUDIO_FILE, 'r', encoding='utf-8') as src_file:       # 打开原始字幕文件
#         src_content = src_file.read()                                            # 读取原始字幕内容
    
#     # ------------
#     # 初始化数据结构
#     # ------------
#     subtitles = []                                                               # 初始化字幕列表
#     src_subtitles = {}                                                           # 初始化原始字幕字典，以序号为键
    
#     # ------------
#     # 解析原始字幕文件
#     # ------------
#     for block in src_content.strip().split('\n\n'):                             # 按空行分割字幕块
#         lines = [line.strip() for line in block.split('\n') if line.strip()]    # 获取非空行列表
#         if len(lines) < 3:                                                       # 如果行数少于3行（序号、时间、内容）
#             continue                                                             # 跳过无效块
        
#         number = int(lines[0])                                                   # 获取字幕序号
#         src_text = ' '.join(lines[2:])                                           # 连接第3行及之后的所有文本内容
#         src_subtitles[number] = src_text                                         # 存储到原始字幕字典中
    
#     # ------------
#     # 解析翻译字幕文件
#     # ------------
#     for block in content.strip().split('\n\n'):                                 # 按空行分割翻译字幕块
#         lines = [line.strip() for line in block.split('\n') if line.strip()]    # 获取非空行列表
#         if len(lines) < 3:                                                       # 如果行数少于3行
#             continue                                                             # 跳过无效块
        
#         try:                                                                     # 尝试解析字幕块
#             number = int(lines[0])                                               # 解析字幕序号
#             start_time, end_time = lines[1].split(' --> ')                       # 分割开始和结束时间
#             start_time = datetime.datetime.strptime(start_time, '%H:%M:%S,%f').time()  # 解析开始时间
#             end_time = datetime.datetime.strptime(end_time, '%H:%M:%S,%f').time()      # 解析结束时间
#             duration = time_diff_seconds(start_time, end_time, datetime.date.today())  # 计算字幕持续时间
#             text = ' '.join(lines[2:])                                           # 连接字幕文本内容
            
#             # ------------
#             # 文本清理处理
#             # ------------
#             text = re.sub(r'\([^)]*\)', '', text).strip()                        # 移除英文圆括号及其内容
#             text = re.sub(r'（[^）]*）', '', text).strip()                        # 移除中文圆括号及其内容
#             text = text.replace('-', '')                                         # 移除横线字符，可继续添加其他可能导致错误的非法字符

#             # ------------
#             # 添加原始字幕文本
#             # ------------
#             origin = src_subtitles.get(number, '')                               # 从原始字幕字典中获取对应的原文

#         except ValueError as e:                                                  # 如果解析出错
#             rprint(Panel(f"Unable to parse subtitle block '{block}', error: {str(e)}, skipping this subtitle block.", title="Error", border_style="red"))  # 打印错误信息
#             continue                                                             # 跳过当前块继续处理
        
#         # ------------
#         # 添加字幕记录
#         # ------------
#         subtitles.append({'number': number, 'start_time': start_time, 'end_time': end_time, 'duration': duration, 'text': text, 'origin': origin})  # 将解析后的字幕信息添加到列表
    
#     # ------------
#     # 创建数据框
#     # ------------
#     df = pd.DataFrame(subtitles)                                                 # 将字幕列表转换为pandas数据框
    
#     # ------------
#     # 字幕时长优化处理
#     # ------------
#     i = 0                                                                        # 初始化循环索引
#     MIN_SUB_DUR = load_key("min_subtitle_duration")                              # 加载最小字幕持续时间配置
#     while i < len(df):                                                           # 遍历所有字幕记录
#         today = datetime.date.today()                                            # 获取今天的日期作为基准日期
#         if df.loc[i, 'duration'] < MIN_SUB_DUR:                                  # 如果当前字幕时长小于最小时长
#             if i < len(df) - 1 and time_diff_seconds(df.loc[i, 'start_time'],df.loc[i+1, 'start_time'],today) < MIN_SUB_DUR:  # 如果不是最后一个字幕且与下一个字幕间隔很短
#                 rprint(f"[bold yellow]Merging subtitles {i+1} and {i+2}[/bold yellow]")  # 打印合并信息
#                 df.loc[i, 'text'] += ' ' + df.loc[i+1, 'text']                   # 合并当前和下一个字幕的翻译文本
#                 df.loc[i, 'origin'] += ' ' + df.loc[i+1, 'origin']               # 合并当前和下一个字幕的原始文本
#                 df.loc[i, 'end_time'] = df.loc[i+1, 'end_time']                  # 将结束时间更新为下一个字幕的结束时间
#                 df.loc[i, 'duration'] = time_diff_seconds(df.loc[i, 'start_time'],df.loc[i, 'end_time'],today)  # 重新计算合并后的持续时间
#                 df = df.drop(i+1).reset_index(drop=True)                         # 删除下一个字幕记录并重置索引
#             else:                                                                # 如果无法合并
#                 if i < len(df) - 1:                                              # 如果不是最后一个字幕
#                     rprint(f"[bold blue]Extending subtitle {i+1} duration to {MIN_SUB_DUR} seconds[/bold blue]")  # 打印延长时间信息
#                     df.loc[i, 'end_time'] = (datetime.datetime.combine(today, df.loc[i, 'start_time']) +   # 计算新的结束时间
#                                             datetime.timedelta(seconds=MIN_SUB_DUR)).time()  # 在开始时间基础上加上最小时长
#                     df.loc[i, 'duration'] = MIN_SUB_DUR                          # 设置持续时间为最小时长
#                 else:                                                            # 如果是最后一个字幕
#                     rprint(f"[bold red]The last subtitle {i+1} duration is less than {MIN_SUB_DUR} seconds, but not extending[/bold red]")  # 打印不延长的警告
#                 i += 1                                                           # 移动到下一个字幕
#         else:                                                                    # 如果时长满足要求
#             i += 1                                                               # 直接移动到下一个字幕
    
#     # ------------
#     # 时间格式转换
#     # ------------
#     df['start_time'] = df['start_time'].apply(lambda x: x.strftime('%H:%M:%S.%f')[:-3])  # 将开始时间转换为字符串格式，保留3位毫秒
#     df['end_time'] = df['end_time'].apply(lambda x: x.strftime('%H:%M:%S.%f')[:-3])      # 将结束时间转换为字符串格式，保留3位毫秒

#     # ------------
#     # 已弃用的二次修剪功能
#     # ------------
#     ##! No longer perform secondary trim
#     # check and trim subtitle length, for twice to ensure the subtitle length is within the limit, 允许tolerance
#     # df['text'] = df.apply(lambda x: check_len_then_trim(x['text'], x['duration']+x['tolerance']), axis=1)

#     return df                                                                    # 返回处理完成的数据框

# ------------
# 音频任务生成主函数
# ------------
@check_file_exists(_AUDIO_TASKS)                                                # 检查音频任务文件是否存在的装饰器
def gen_audio_task_main():
    """
    生成音频合成任务的主函数。
    
    该函数从翻译结果文件中读取已翻译的文本片段，为每个片段创建对应的
    音频合成任务，并将任务列表保存为JSON格式的配置文件。
    
    处理流程：
    1. 读取翻译结果文件 (_TRANSLATE_RESULT)
    2. 为每个翻译片段创建音频任务对象
    3. 生成对应的输出文件名（基于片段ID）
    4. 将任务列表保存到音频任务文件 (_AUDIO_TASKS)
    5. 在控制台显示处理结果统计
    
    任务对象结构：
        {
            "task_id": int,           # 任务ID，对应翻译片段ID
            "text_to_synthesize": str, # 待合成的文本内容
            "output_filename": str     # 输出音频文件名（如"1.mp3"）
        }
    
    Files:
        Input: _TRANSLATE_RESULT (翻译结果JSON文件)
        Output: _AUDIO_TASKS (音频任务JSON文件)
    
    Raises:
        FileNotFoundError: 当翻译结果文件不存在时
        json.JSONDecodeError: 当翻译结果文件格式无效时
        PermissionError: 当无法写入输出目录时
    
    Notes:
        - 使用@check_file_exists装饰器确保不会覆盖已存在的音频任务文件
        - 自动创建output目录如果不存在
        - JSON输出使用UTF-8编码并格式化缩进以便阅读
    """
    # 读取翻译结果
    with open(_TRANSLATE_RESULT, 'r', encoding='utf-8') as f:
        translate_result = json.load(f)
    
    # 为每个翻译片段生成音频任务
    audio_tasks = []
    for segment in translate_result:
        task = {
            "task_id": segment['id'],
            "text_to_synthesize": segment['translated_text'],
            "output_filename": f"{segment['id']}.mp3"
        }
        audio_tasks.append(task)
    
    # 保存音频任务到JSON文件
    os.makedirs('output', exist_ok=True)
    with open(_AUDIO_TASKS, 'w', encoding='utf-8') as f:
        json.dump(audio_tasks, f, ensure_ascii=False, indent=2)
    
    console.print(f"📊 Generated {len(audio_tasks)} audio tasks")
    rprint(Panel(f"Successfully generated {_AUDIO_TASKS}", title="Success", border_style="green"))  # 打印成功信息

# ------------
# 主程序入口
# ------------
if __name__ == '__main__':                                                      # 如果直接运行此脚本
    gen_audio_task_main()                                                        # 执行音频任务生成主函数