import gradio as gr

from main import calculate_cer_both_versions, read_srt_text
from visualize import generate_html_report


def process_srt_files(reference_file, hypothesis_file):
    try:
        # Handle both file objects and string paths
        ref_path = reference_file.name if hasattr(
            reference_file, 'name') else reference_file
        hyp_path = hypothesis_file.name if hasattr(
            hypothesis_file, 'name') else hypothesis_file

        reference_text = read_srt_text(ref_path)
        hypothesis_text = read_srt_text(hyp_path)

        metrics_no_punct, metrics_with_punct = calculate_cer_both_versions(
            ref_path, hyp_path)

        html_content = generate_html_report(
            reference_text, hypothesis_text, metrics_no_punct, metrics_with_punct)

        return html_content
    except Exception as e:
        return f"An error occurred: {str(e)}"


def load_example():
    return "gold.srt", "bert.srt"


with gr.Blocks() as iface:
    gr.Markdown("# SRT File Comparison and CER Calculation")
    gr.Markdown(
        "## Please upload the golden reference SRT and the target SRT for calculating the CER.")
    gr.Markdown(        
        "Note: Only CER is supported at the moment, WER will be added in a future version.")

    with gr.Row():
        ref_file = gr.File(label="Reference (Golden) SRT File")
        hyp_file = gr.File(label="Target SRT File")

    with gr.Row():
        example_btn = gr.Button("Load Example")
        process_btn = gr.Button("Get CER", variant="primary")

    output = gr.HTML(label="Results")

    process_btn.click(
        fn=process_srt_files,
        inputs=[ref_file, hyp_file],
        outputs=output
    )

    example_btn.click(
        fn=load_example,
        inputs=None,
        outputs=[ref_file, hyp_file]
    )

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