import gradio as gr
import cv2
from paddleocr import PaddleOCR
import datetime
import difflib
import os

# Khởi tạo mô hình OCR tiếng Việt với các tham số tương thích phiên bản mới
ocr = PaddleOCR(lang='vi')

def format_timestamp(ms):
    """Chuyển đổi mili-giây sang định dạng chuẩn SRT: HH:MM:SS,mmm"""
    seconds, milliseconds = divmod(ms, 1000)
    minutes, seconds = divmod(seconds, 60)
    hours, minutes = divmod(minutes, 60)
    return f"{int(hours):02d}:{int(minutes):02d}:{int(seconds):02d},{int(milliseconds):03d}"

def is_similar(text1, text2, threshold=0.85):
    """So sánh độ giống nhau giữa 2 chuỗi để gộp khung hình (tránh nhiễu OCR)"""
    return difflib.SequenceMatcher(None, text1, text2).ratio() > threshold

def process_video(video_path, progress=gr.Progress()):
    if not video_path:
        return None, "Vui lòng tải lên một video."

    cap = cv2.VideoCapture(video_path)
    
    # Lấy thông số video
    fps = cap.get(cv2.CAP_PROP_FPS)
    total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
    
    # Thiết lập quét 2 khung hình mỗi giây để tối ưu tốc độ
    frames_to_skip = int(fps / 2) if fps > 0 else 15 
    interval_ms = (frames_to_skip / fps) * 1000 if fps > 0 else 500

    srt_blocks = []
    current_text = ""
    start_time_ms = 0
    last_seen_ms = 0
    
    frame_count = 0
    
    while cap.isOpened():
        ret, frame = cap.read()
        if not ret:
            break
            
        # Chỉ quét các khung hình được chỉ định
        if frame_count % frames_to_skip == 0:
            current_ms = cap.get(cv2.CAP_PROP_POS_MSEC)
            
            # Báo cáo tiến độ lên giao diện Gradio
            progress(frame_count / total_frames, desc=f"Đang xử lý thời gian: {format_timestamp(current_ms)}")
            
            # TỐI ƯU HÓA: Chỉ cắt lấy 30% khu vực bên dưới của video (nơi đặt phụ đề)
            h, w, _ = frame.shape
            cropped_frame = frame[int(h * 0.7):h, 0:w]
            
            # Chạy OCR trên ảnh đã cắt
            result = ocr.ocr(cropped_frame, cls=False)
            
            text_lines = []
            if result and result[0]:
                text_lines = [line[1][0] for line in result[0]]
            
            frame_text = " ".join(text_lines).strip()
            
            # THUẬT TOÁN GỘP TIMECODE
            if frame_text:
                if not current_text:
                    # Bắt đầu một câu phụ đề mới
                    current_text = frame_text
                    start_time_ms = current_ms
                    last_seen_ms = current_ms
                elif is_similar(current_text, frame_text):
                    # Kéo dài thời gian hiển thị nếu câu phụ đề vẫn giống nhau
                    last_seen_ms = current_ms
                else:
                    # Phụ đề đã thay đổi -> Lưu câu cũ vào danh sách SRT
                    srt_blocks.append((start_time_ms, last_seen_ms + interval_ms, current_text))
                    # Cập nhật câu mới
                    current_text = frame_text
                    start_time_ms = current_ms
                    last_seen_ms = current_ms
            else:
                # Không có text -> Lưu câu cũ nếu đang có
                if current_text:
                    srt_blocks.append((start_time_ms, last_seen_ms + interval_ms, current_text))
                    current_text = ""
                    
        frame_count += 1

    # Lưu câu cuối cùng nếu video kết thúc đột ngột
    if current_text:
        srt_blocks.append((start_time_ms, last_seen_ms + interval_ms, current_text))

    cap.release()

    # Định dạng thành cấu trúc file SRT chuẩn
    srt_content = ""
    for i, (start_ms, end_ms, text) in enumerate(srt_blocks, 1):
        srt_content += f"{i}\n"
        srt_content += f"{format_timestamp(start_ms)} --> {format_timestamp(end_ms)}\n"
        srt_content += f"{text}\n\n"

    # Xuất file
    output_filename = "output_subtitles.srt"
    with open(output_filename, "w", encoding="utf-8") as f:
        f.write(srt_content)

    return output_filename, f"Hoàn tất! Trích xuất được {len(srt_blocks)} dòng phụ đề."

# Giao diện Gradio
with gr.Blocks(title="Video OCR to SRT") as demo:
    gr.Markdown("# Trích xuất Phụ đề Video (OCR -> SRT)")
    
    with gr.Row():
        with gr.Column():
            video_input = gr.Video(label="Tải video lên")
            process_btn = gr.Button("Bắt đầu trích xuất", variant="primary")
        
        with gr.Column():
            file_output = gr.File(label="Tải xuống file SRT")
            status_output = gr.Textbox(label="Trạng thái", interactive=False)
            
    process_btn.click(
        fn=process_video,
        inputs=[video_input],
        outputs=[file_output, status_output]
    )

if __name__ == "__main__":
    demo.launch()