import cv2
import gradio as gr
from paddleocr import PaddleOCR
import os
import gc

# Time Formatter (SRT Standard)
def format_srt_time(seconds):
    hours = int(seconds // 3600)
    minutes = int((seconds % 3600) // 60)
    secs = int(seconds % 60)
    millis = int((seconds % 1) * 1000)
    return f"{hours:02d}:{minutes:02d}:{secs:02d},{millis:03d}"

def extract_srt(video_path, lang_choice):
    if not video_path:
        return None
    
    lang_codes = {
        "Chinese (S)": "ch", "Chinese (T)": "chinese_cht", 
        "English": "en", "Japanese": "japan", 
        "Korean": "korean", "Thai": "th"
    }
    
    # Error ကင်းစေရန် Parameter အသစ်များဖြင့် ပြင်ဆင်ထားသည်
    try:
        ocr = PaddleOCR(use_angle_cls=True, lang=lang_codes[lang_choice], use_gpu=False)
    except Exception as e:
        return None

    cap = cv2.VideoCapture(video_path)
    fps = cap.get(cv2.CAP_PROP_FPS)
    if fps <= 0: fps = 24
    
    subs = []
    current_text = ""
    start_time = 0
    frame_idx = 0
    
    # ၁ စက္ကန့်လျှင် ၄ ကြိမ်နှုန်းဖြင့် စာသားကို အသေးစိတ်ဖတ်မည်
    interval = max(1, int(fps / 4))

    while cap.isOpened():
        ret, frame = cap.read()
        if not ret: break
        
        if frame_idx % interval == 0:
            timestamp = frame_idx / fps
            
            result = ocr.ocr(frame, cls=True)
            frame_text = ""
            
            if result and result[0]:
                lines = [line[1][0] for line in result[0]]
                frame_text = " ".join(lines).strip()
            
            # စာသားအပြောင်းအလဲပေါ်မူတည်၍ Timer သတ်မှတ်ခြင်း
            if frame_text != current_text:
                if current_text:
                    subs.append({
                        "start": format_srt_time(start_time),
                        "end": format_srt_time(timestamp),
                        "text": current_text
                    })
                
                if frame_text:
                    start_time = timestamp
                    current_text = frame_text
                else:
                    current_text = ""
            
            del frame
            gc.collect()
                    
        frame_idx += 1
    
    cap.release()

    # နောက်ဆုံးစာကြောင်းအတွက် သိမ်းဆည်းခြင်း
    if current_text:
        subs.append({"start": format_srt_time(start_time), "end": format_srt_time(frame_idx/fps), "text": current_text})

    # SRT File အဖြစ် ထုတ်ပေးခြင်း
    srt_path = "output_subtitles.srt"
    with open(srt_path, "w", encoding="utf-8") as f:
        for i, s in enumerate(subs):
            f.write(f"{i+1}\n{s['start']} --> {s['end']}\n{s['text']}\n\n")
            
    return srt_path

# UI Logic
with gr.Blocks() as demo:
    gr.Markdown("### 🎬 Video Hardsub to SRT (No-Error Version)")
    with gr.Row():
        with gr.Column():
            video_input = gr.Video()
            lang_dropdown = gr.Dropdown(
                choices=["Chinese (S)", "Chinese (T)", "English", "Japanese", "Korean", "Thai"], 
                value="Chinese (S)", 
                label="ဘာသာစကားရွေးပါ"
            )
            process_btn = gr.Button("SRT ဖိုင်ထုတ်ယူမည်", variant="primary")
        with gr.Column():
            file_output = gr.File(label="Download Subtitles (.srt)")

    process_btn.click(extract_srt, inputs=[video_input, lang_dropdown], outputs=file_output)

demo.launch()