import traceback
import gradio as gr
from utils.get_RGB_image import get_RGB_image, is_online_file, steam_online_file
from pdf2image import convert_from_path, convert_from_bytes

def similarity_fn(document_image_1, document_image_2):
    xxx = 'BOTH' if document_image_1 and document_image_2 else 'ONE' if document_image_1 or document_image_2 else 'NONE'
    return gr.HTML(f'<pre>Similarity between the two documents: {xxx}<pre>', visible=True)

def load_image(filename, page = 0):
    try:
        image = None
        try:
            if (is_online_file(filename)):
                image = get_RGB_image(convert_from_bytes(steam_online_file(filename))[page])
            else:
                image = get_RGB_image(convert_from_path(filename)[page])
        except:
            image = get_RGB_image(filename)
        return [
            gr.Image(value=image, visible=True), 
            None
        ]
    except:
        error = traceback.format_exc()
        return [None, gr.HTML(value=error, visible=True)]
    
def preview_url(url, page = 0):
    [image, error] = load_image(url, page = page)
    if image:
        return [gr.Tabs(selected=0), image, error]
    else:
        return [gr.Tabs(selected=1), image, error]
    
def document_view(document_number: int):
    gr.HTML(value=f'<h4>Load the {"first" if document_number == 1 else "second"} PDF or Document Image<h4>', elem_classes=['center'])
    with gr.Tabs() as document_tabs:
        with gr.Tab("From Image", id=0):
            document = gr.Image(type="pil", label=f"Document {document_number}", visible=False)
            document_error_message = gr.HTML(label="Error Message", visible=False)
            document_preview = gr.UploadButton(
                "Click to PDF or Document Image", 
                file_types=["image", ".pdf"], 
                file_count="single")
        with gr.Tab("From URL", id=1):
            document_url = gr.Textbox(
                label=f"Document {document_number} URL",
                info="Paste a Link/URL to PDF or Document Image",
                placeholder="https://datasets-server.huggingface.co/.../image.jpg")
            document_url_error_message = gr.HTML(label="Error Message", visible=False)
            document_url_preview = gr.Button(value="Preview", variant="primary")
    document_preview.upload(
         fn = lambda file: load_image(file.name), 
         inputs = [document_preview], 
         outputs = [document, document_error_message])
    document_url_preview.click(
        fn = preview_url, 
        inputs = [document_url], 
        outputs = [document_tabs, document, document_url_error_message])
    return document

def app():
    title = 'Document Similarity Search Using Visual Layout Features'
    description = f"<h2>{title}<h2>"
    css = '''
    image { max-height="86vh" !important; }
    .center { display: flex; flex: 1 1 auto; align-items: center; align-content: center; justify-content: center; justify-items: center; }
    .hr { width: 100%; display: block; padding: 0; margin: 0; background: gray; height: 4px;  border: none; }
  '''
    with gr.Blocks(title=title, css=css) as app:
        with gr.Row():
            gr.HTML(value=description, elem_classes=['center'])
        with gr.Row(equal_height = False):
            with gr.Column():
                document_1_image = document_view(1)
            with gr.Column():
                document_2_image = document_view(2)
        gr.HTML('<hr/>', elem_classes=['hr'])
        with gr.Row(elem_classes=['center']):
            with gr.Column():
                submit = gr.Button(value="Similarity", variant="primary")
                reset = gr.Button(value="Reset", variant="secondary")
            with gr.Column():
                similarity_output = gr.HTML(visible=False)
            submit.click(
                fn=similarity_fn,
                inputs=[document_1_image, document_2_image],
                outputs=[similarity_output])
    return app.launch(debug=True)