import gradio as gr
import os
from videocr import save_subtitles_to_file

# Define OCR function
def run_video_ocr(input_video, output_file_name, language_code, use_gpu, start_time, end_time, confidence_threshold, similarity_threshold, frames_to_skip, crop_x, crop_y, crop_width, crop_height):
    try:
        # Ensure the output directory exists
        persistent_dir = '/persistent'
        if not os.path.exists(persistent_dir):
            os.makedirs(persistent_dir)
        
        # Define full path for the output file
        output_path = os.path.join(persistent_dir, output_file_name)

        # Save the subtitles to file
        save_subtitles_to_file(
            input_video.name,  # The uploaded video file
            output_path,       # Output .srt file path
            lang=language_code,
            use_gpu=use_gpu,
            time_start=start_time,
            time_end=end_time,
            conf_threshold=confidence_threshold,
            sim_threshold=similarity_threshold,
            frames_to_skip=frames_to_skip,
            crop_x=crop_x,
            crop_y=crop_y,
            crop_width=crop_width,
            crop_height=crop_height
        )

        return f"Subtitle extraction completed! File saved to {output_path}"
    except Exception as e:
        return f"Error: {str(e)}"

# Define Gradio interface
def video_ocr_interface():
    with gr.Blocks() as demo:
        with gr.Row():
            input_video = gr.File(label="Upload Video", type="file")
            output_file_name = gr.Textbox(label="Output File Name (.srt)", value="subtitle.srt")
            language_code = gr.Textbox(label="Language Code", value="ch")
            use_gpu = gr.Checkbox(label="Use GPU", value=True)
        
        with gr.Row():
            start_time = gr.Textbox(label="Start Time (HH:MM:SS)", value="00:00:00")
            end_time = gr.Textbox(label="End Time (HH:MM:SS)", value="")
        
        with gr.Row():
            confidence_threshold = gr.Slider(label="Confidence Threshold", minimum=0, maximum=100, value=75)
            similarity_threshold = gr.Slider(label="Similarity Threshold", minimum=0, maximum=100, value=80)
        
        with gr.Row():
            frames_to_skip = gr.Slider(label="Frames to Skip", minimum=0, maximum=10, value=0)
            crop_x = gr.Number(label="Crop X", value=0)
            crop_y = gr.Number(label="Crop Y", value=0)
            crop_width = gr.Number(label="Crop Width", value=0)
            crop_height = gr.Number(label="Crop Height", value=0)

        submit_btn = gr.Button("Start OCR")

        output_label = gr.Textbox(label="Status", interactive=False)

        # Define what happens when the button is clicked
        submit_btn.click(
            fn=run_video_ocr,
            inputs=[input_video, output_file_name, language_code, use_gpu, start_time, end_time, confidence_threshold, similarity_threshold, frames_to_skip, crop_x, crop_y, crop_width, crop_height],
            outputs=output_label
        )
    
    return demo

# Launch the Gradio interface
demo = video_ocr_interface()
demo.launch()