import os
import torch
import argparse
import gradio as gr
import spaces

from utils import Image2Text

@spaces.GPU(duration=60)
def predict(input_img):
    global image_to_text
    contents = image_to_text.get_text(input_img, num_beams=4)
    return contents

def get_image_path_list(folder_name):
    image_basename_list = os.listdir(folder_name)
    image_path_list = sorted([os.path.join(folder_name, basename) for basename in image_basename_list])
    return image_path_list

def swap_to_gallery(images):
    return gr.update(value=images, visible=True)

def swap_to_gallery(images):
    return gr.update(value=images, visible=True), gr.update(visible=True), gr.update(visible=False)

def remove_back_to_files():
    return gr.update(visible=False), gr.update(visible=False), gr.update(visible=True)

def upload_example_to_gallery(images):
    return gr.update(value=images, visible=True)
    
def get_example():
    return ['./samples/sample1.jpg', './samples/sample2.jpg']

examples_path = os.path.dirname(__file__)


title = r"""
<h1 align="center">🍩 for Hwp math problems</h1>
<div align="center">
    <h3>
        I have launched a service that includes HWP automation.<br>
        Automatically convert the processed images into HWP documents. <a href="https://hwpmath.duckdns.org/">[Try now! 🤗]</a>
    </h3>
</div>
"""

description = r"""
<b>Official 🤗 Gradio demo</b> for <a href='https://github.com/sooooner/DonutMathHWP' target='_blank'><b>DonutMathHWP</b></a>.<br>
<br>
❗️[<b>Important</b>] HWP Conversion Steps:<br>
1️⃣ <b>Image Upload</b>: Upload an image containing a math problem or solution, including any mathematical formulas. Note that text recognition for problems is not included, so please ensure the math content <b>occupies the majority of the image</b>.<br>
2️⃣ <b>Start Conversion</b>: Click the <b>Submit</b> button to initiate the conversion.<br>
3️⃣ <b>Conversion Output</b>: In the converted output, <b>text enclosed within $ symbols</b> represents mathematical formulas. When using the output in an HWP document, treat all segments surrounded by $ symbols as formulas for accurate formatting.<br>
"""

example = r"""
<div>
<h3>💡 <b>Example Images:</b> Sample Images of Math Problems and Solutions</h3>

<b>Problem:</b> The image should include a math problem with relevant formulas or symbols, taking up most of the image area for clear recognition.  
<b>Solution:</b> The image should show solution steps with formulas and calculations, occupying most of the space and in sequential order.  
</div>
"""

article = r"""<br>
📧 <b>Contact</b>
<br>
If you have any questions, please feel free to reach me out at <b>tnsgh0101@gmail.com</b>
"""

css = """
    .gradio-container {
        width: 85% !important
    }
    .highlighted-examples {
        border: 4px solid #FFD700;
        padding: 20px;
        margin: 30px 0;
        background-color: #FFFACD;
        border-radius: 15px;
        font-size: 1.2em;
    }
    .highlighted-examples:hover {
        box-shadow: 0 0 25px rgba(255, 215, 0, 0.9);
        cursor: pointer;
        transform: scale(1.02);
    }
"""

if __name__ == "__main__":
    repo_id = os.getenv('MODEL_REPO_ID')
    hf_token = os.environ.get("HF_TOKEN")

    device = "cuda"
    image_to_text = Image2Text(repo_id, hf_token=hf_token, device=device)

    with gr.Blocks(css=css) as demo:
        gr.Markdown(title)
        gr.Markdown(description)
        with gr.Row():
            with gr.Column():
                files = gr.File(
                            label="Drag (Select) 1 or more photos",
                            file_types=["image"],
                            file_count="multiple",
                            visible=True
                        )
                uploaded_files = gr.Gallery(show_label=False, visible=False, columns=2)
                with gr.Column(visible=False) as clear_button:
                    remove_and_reupload = gr.ClearButton(value="Remove and upload new ones", components=files, size="sm")
                submit = gr.Button("Submit")

            with gr.Column():
                extracted_texts_output = gr.JSON(label="Extracted Texts")

            files.upload(fn=swap_to_gallery, inputs=files, outputs=[uploaded_files, clear_button, files])
            remove_and_reupload.click(fn=remove_back_to_files, outputs=[uploaded_files, clear_button, files])

            submit.click(
                fn=predict,
                inputs=files,
                outputs=extracted_texts_output
            )

        # gr.Examples(
        #     examples=[[get_image_path_list('./samples/problems')], [get_image_path_list('./samples/solutions')]],
        #     # examples=[
        #     #     (get_image_path_list('./samples/problems'), "Problems"),
        #     #     (get_image_path_list('./samples/solutions'), "Solutions")
        #     # ],
        #     fn=swap_to_gallery,
        #     run_on_click=True,
        #     inputs=files,
        #     outputs=[uploaded_files, clear_button, files],
        #     cache_examples=True,
        #     label="💡 Example Images",
        #     elem_id="highlighted-examples"
        # )

        gr.Markdown(example)    
        gr.Examples(
            examples=[[get_image_path_list('./samples/problems')], [get_image_path_list('./samples/solutions')]],
            inputs=files,
            run_on_click=True,
            outputs=[uploaded_files, clear_button, files],
            fn=swap_to_gallery,
            cache_examples=True,
            label=r"Select a sample images of the problem or solution",
            elem_id="highlighted-examples"
        )
        
        gr.Markdown(article)
        
    demo.launch()