# core.py
import google.generativeai as genai
import pysrt
import re
import os

# 解析 SRT 檔案
def parse_srt(srt_path):
    try:
        subs = pysrt.open(srt_path, encoding='utf-8')
        srt_data = {}
        for sub in subs:
            srt_data[sub.index] = {"time": str(sub.start) + " --> " + str(sub.end), "text": sub.text}
        return srt_data
    except Exception as e:
        print(f"解析字幕檔案錯誤: {e}")
        return None

# 預處理逐字稿
def preprocess_transcript(transcript):
        # 移除 SSML 標籤
    transcript = re.sub(r'<[^>]*?>', '', transcript)
    # 移除多個連續空格
    transcript = re.sub(r'\s+', ' ', transcript)
    # 移除所有換行符號
    transcript = transcript.replace('\n','')
    # 移除文字前後的空白
    transcript = transcript.strip()
    return transcript

def process_files(api_key, test_transcript_file, test_srt_file, batch_size):
    # 1. 配置 Gemini API
    try:
        genai.configure(api_key=api_key)
        model = genai.GenerativeModel('gemini-2.0-flash-exp') # 測試金鑰是否正確
    except Exception as e:
        return f"API 金鑰錯誤，請檢查設定: {e}", None, None

    # 檢查測試用逐字稿檔案是否存在
    if not os.path.exists(test_transcript_file):
        return f"錯誤: 找不到逐字稿檔案: {test_transcript_file}，請確認您已上傳該檔案", None, None
    else:
        with open(test_transcript_file, 'r', encoding='utf-8') as f:
            transcript_content = f.read()
        transcript_content = preprocess_transcript(transcript_content)

    # 檢查測試用 srt 檔案是否存在
    if not os.path.exists(test_srt_file):
        return f"錯誤: 找不到口語稿檔案: {test_srt_file}，請確認您已上傳該檔案", None, None
    else:
        srt_data = parse_srt(test_srt_file)
        if srt_data is None:
            return "錯誤: 解析口語稿檔案失敗，請確認檔案內容", None, None
    
    # 儲存所有報告
    all_reports = []
    keys = list(srt_data.keys()) # 取得編號

    # 處理每個批次
    for i in range(0, len(keys), batch_size):
        batch_keys = keys[i:i + batch_size]
        batch_data = {key: srt_data[key] for key in batch_keys} # 取得該批次的資料
        subtitle_content = "\n".join([f"編號{index}:{data['text']}" for index, data in batch_data.items()])
        prompt = (
                "請扮演一位專業的文字編輯，比對以下字幕檔和逐字稿，並以逐字稿的內容為準，修正字幕檔中因語音轉文字造成的錯別字、漏字、多字、詞語誤用等問題，請勿修改語氣和修辭，必須嚴格保持原始字幕的時間戳和斷句結構，不要拆分或合併任何段落，不要新增或移除任何字幕段落，直接輸出修正後的字幕檔。\n"
                "接著，請將**當前批次**所有修改的內容整理成一個修改清單，並放在所有字幕內容之後。\n"
                "請使用以下格式輸出修改清單： \n"
                "編號: 原文: 修正後\n"
                "請使用 `<<<分隔符號>>>` 來分隔字幕與修改清單。\n"
                 "請按照以下步驟進行：\n"
                "1. 比對字幕檔和逐字稿。\n"
                "2. 以逐字稿的內容為準，修正字幕檔中因語音轉文字造成的錯別字、漏字、多字和詞語誤用、以及標點符號漏缺等問題。\n"
                "3. 必須嚴格保持原始字幕的時間戳和斷句結構，不要拆分或合併任何段落，不要新增或移除任何字幕段落，直接輸出修正後的字幕檔。\n"
                "4. 讀取整份逐字稿時，請忽略SSML標籤，直接比對文字與標點符號，請勿修改任何語義。\n"
                 "字幕檔內容：\n"
                f"{subtitle_content}\n\n"
                "逐字稿內容：\n"
                 f"{transcript_content}\n\n"
               )
        try:
            response = model.generate_content(prompt)
            corrected_subtitle = response.text
            print(f"第 {i // batch_size + 1} 批次 Gemini 模型的回應：")
            print(corrected_subtitle)

            # 使用 re.split 分割字幕和報告
            parts = re.split(r'<<<分隔符號>>>', corrected_subtitle, maxsplit=1) # 只分割一次
            corrected_lines = parts[0].split('\n')  # 取得字幕文字
            report = ""
            if len(parts) > 1:
                report = parts[1] # 取得修改清單文字
                report_lines = report.strip().split('\n')
                all_reports.extend(report_lines)

            for line in corrected_lines:
                match = re.match(r'編號(\d+):(.*)', line)
                if match:
                   index = int(match.group(1))
                   corrected_text = match.group(2).strip()
                   srt_data[index]['text'] = corrected_text
        except Exception as e:
            return f"第 {i // batch_size + 1} 批次修正過程失敗：{e}", None, None
    
    # 寫回 SRT 檔案
    new_subs = []
    for index, data in srt_data.items():
       time_str = data['time']
       start_str, end_str = time_str.split(" --> ")
       start = pysrt.srtitem.SubRipTime.from_string(start_str)
       end = pysrt.srtitem.SubRipTime.from_string(end_str)
       new_subs.append(pysrt.SubRipItem(index=index, start=start, end=end, text=data['text']))
    
    # 將 new_subs 轉換為 SubRipFile 物件
    new_srt = pysrt.SubRipFile(items=new_subs)
    
    return None, new_srt, all_reports