import gradio as gr
from PIL import Image

HEAD_HTML = """
<link href='https://fonts.googleapis.com/css?family=PT Mono' rel='stylesheet'>
"""
DEFAULT_CASE = "case_2"
CASE_LIST = ["case_0", "case_1", "case_2"]
data = {}

for case_name in CASE_LIST:
    with open(f"resources/{case_name}/prompt_0.txt", "r") as f:
        data[case_name] = {
            "prompt_0": f.read(),
        }
    with open(f"resources/{case_name}/output_0.txt", "r") as f:
        data[case_name]["output_0"] = f.read()
        data[case_name]["output_1"] = f.read()

    with open(f"resources/{case_name}/prompt_1.txt", "r") as f:
        data[case_name]["prompt_1"] = f.read()
    with open(f"resources/{case_name}/output_1.txt", "r") as f:
        data[case_name]["output_1"] = f.read()
        
    demo_image = Image.open(
        f"resources/{case_name}/pdf_page.png"
    )
    data[case_name]["image"] = demo_image

with gr.Blocks(
    theme=gr.themes.Ocean(
        font_mono="PT Mono",
    ),
    head=HEAD_HTML,
) as demo:
    gr.Markdown("# 2D Layout-Preserving Text Serialization Demo")
    with gr.Row():
        gr.Markdown("## Input Image")
        case_name = gr.Dropdown(
            choices=CASE_LIST,
            label="Select a case",
            value=DEFAULT_CASE,
        )
        
    demo_image = gr.Image(
        data[DEFAULT_CASE]["image"],
        show_label=False,
        container=False,
        height=700,
        width=600,
    )
    gr.Markdown("## Sample Prompt & Response from LLM")
    with gr.Tabs():
        with gr.TabItem("Prompt"):
            with gr.Row():
                prompt_0 = gr.Code(
                    data[DEFAULT_CASE]["prompt_0"], 
                    language="markdown", 
                    show_line_numbers=False, 
                    interactive=False,
                    label="2D Layout-Preserving Text"
                )
                prompt_1 = gr.Code(
                    data[DEFAULT_CASE]["prompt_1"],
                    language="markdown", 
                    show_line_numbers=False, 
                    interactive=False,
                    label="1D Text"
                )
                
        with gr.TabItem("Response"):
            with gr.Row():
                output_0 = gr.Markdown(
                    data[DEFAULT_CASE]["output_0"], 
                    container=False
                )
                output_1 = gr.Markdown(
                    data[DEFAULT_CASE]["output_1"], 
                    container=False
                )
    
    case_name.change(
        lambda case: [
            data[case]["image"],
            data[case]["prompt_0"],
            data[case]["prompt_1"],
            data[case]["output_0"],
            data[case]["output_1"],
        ],
        inputs=case_name,
        outputs=[
            demo_image,
            prompt_0,
            prompt_1,
            output_0,
            output_1,
        ],
    )
    demo.queue().launch()