import os
import gradio as gr
import requests
import json
import io
from gradio.components import Image
from PIL import Image as PILImage, ImageDraw, ImageFont  # This import may be needed if you're processing images

from PIL import Image

def process_image(image1, image2):
    try:
        img_bytes1 = io.BytesIO()
        image1.save(img_bytes1, format="JPEG")
        img_bytes1.seek(0)
    except:
        return ["Failed to open image", {"resultCode": "Error", "result": "Failed to open image"}]

    try:
        img_bytes2 = io.BytesIO()
        image2.save(img_bytes2, format="JPEG")
        img_bytes2.seek(0)
    except:
        pass   
    
    url = "http://127.0.0.1:9000/process_image"
    files = {'image1': img_bytes1, 'image2': img_bytes2}
    result = requests.post(url=url, files=files)
    if result.ok:
        json_result = result.json()
        if json_result.get("resultCode") == "Error":
            return ["Failed to process image", {"resultCode": "Error", "result": "Failed to process image"}]

        html = None    
        images = None
        mrz = None
        barcode = None

        r = json_result.get("result", {})
        table_value = ""

        if r.get('MRZ') is not None:
            mrz = r.get('MRZ')

        if r.get('Barcode') is not None:
            barcode = r.get('Barcode')

        for key, value in r.items():
            if key == 'Images' or key == 'MRZ' or key == 'Barcode':
                continue

            mrz_value = ''
            if mrz is not None and mrz.get(key) is not None:
                mrz_value = mrz[key]
                del mrz[key]

            barcode_value = ''
            if barcode is not None and barcode.get(key) is not None:
                barcode_value = barcode[key]
                del barcode[key]

            if mrz is not None and barcode is not None:
                row_value = ("<tr>"
                                "<td>{key}</td>"
                                "<td>{value}</td>"
                                "<td>{mrz_value}</td>"
                                "<td>{barcode_value}</td>"
                            "</tr>".format(key=key, value=value, mrz_value=mrz_value, barcode_value=barcode_value))
            elif mrz is not None:
                row_value = ("<tr>"
                                "<td>{key}</td>"
                                "<td>{value}</td>"
                                "<td>{mrz_value}</td>"
                            "</tr>".format(key=key, value=value, mrz_value=mrz_value))
            elif barcode is not None:    
                row_value = ("<tr>"
                                "<td>{key}</td>"
                                "<td>{value}</td>"
                                "<td>{barcode_value}</td>"
                            "</tr>".format(key=key, value=value, barcode_value=barcode_value))
            else:
                row_value = ("<tr>"
                                "<td>{key}</td>"
                                "<td>{value}</td>"
                            "</tr>".format(key=key, value=value))

            table_value = table_value + row_value

        if mrz is not None and barcode is not None:
            for key, value in mrz.items():
                if key == 'MRZ':
                    value = value.replace('<', '&lt;')
                    value = value.replace(',', '<p>')

                row_value = ("<tr>"
                                "<td>{key}</td>"
                                "<td>{value}</td>"
                                "<td>{mrz_value}</td>"
                                "<td></td>"
                            "</tr>".format(key=key, value='', mrz_value=value))
                table_value = table_value + row_value

            for key, value in barcode.items():
                if key == 'Barcode':
                    value = value.replace('<', '&lt;')
                    value = value.replace(',', '<p>')

                row_value = ("<tr>"
                                "<td>{key}</td>"
                                "<td>{value}</td>"
                                "<td></td>"
                                "<td>{barcode_value}</td>"
                            "</tr>".format(key=key, value='', barcode_value=value))
                table_value = table_value + row_value
        elif mrz is not None:
            for key, value in mrz.items():
                if key == 'MRZ':
                    value = value.replace('<', '&lt;')
                    value = value.replace(',', '<p>')

                row_value = ("<tr>"
                                "<td>{key}</td>"
                                "<td>{value}</td>"
                                "<td>{mrz_value}</td>"
                            "</tr>".format(key=key, value='', mrz_value=value))
                table_value = table_value + row_value                
        elif barcode is not None:
            for key, value in barcode.items():
                if key == 'Barcode':
                    value = value.replace('<', '&lt;')
                    value = value.replace(',', '<p>')

                row_value = ("<tr>"
                                "<td>{key}</td>"
                                "<td>{value}</td>"
                                "<td>{barcode_value}</td>"
                            "</tr>".format(key=key, value='', barcode_value=value))
                table_value = table_value + row_value

        if mrz is not None and barcode is not None:
            html = ("<table>"
                        "<tr>"
                            "<th>Field</th>"
                            "<th>OCR</th>"
                            "<th>MRZ</th>"
                            "<th>Barcode</th>"
                        "</tr>"
                        "{table_value}"
                        "</table>".format(table_value=table_value))
        elif mrz is not None:
            html = ("<table>"
                        "<tr>"
                            "<th style=""width:20%"">Field</th>"
                            "<th style=""width:40%"">OCR</th>"
                            "<th style=""width:40%"">MRZ</th>"
                        "</tr>"
                        "{table_value}"
                        "</table>".format(table_value=table_value))
        elif barcode is not None:
            html = ("<table>"
                        "<tr>"
                            "<th style=""width:20%"">Field</th>"
                            "<th style=""width:40%"">OCR</th>"
                            "<th style=""width:40%"">Barcode</th>"
                        "</tr>"
                        "{table_value}"
                        "</table>".format(table_value=table_value))
        else:
            html = ("<table>"
                        "<tr>"
                            "<th>Field</th>"
                            "<th>OCR</th>"
                        "</tr>"
                        "{table_value}"
                        "</table>".format(table_value=table_value))

        
        table_value = ""
        for key, value in r.items():
            if key == 'Images':
                for image_key, image_value in value.items():
                    row_value = ("<tr>"
                                    "<td>{key}</td>"
                                    "<td><img src=""data:image/png;base64,{base64_image} width = '200'  height= '100' /></td>"
                                "</tr>".format(key=image_key, base64_image=image_value))
                    table_value = table_value + row_value

        images = ("<table>"
                    "<tr>"
                        "<th>Field</th>"
                        "<th>Image</th>"
                    "</tr>"
                    "{table_value}"
                    "</table>".format(table_value=table_value))

        return [html + "<br/><br/>" +  images, json_result]
    else:
        return [result.text, {"resultCode": "Error", "result": result.text}]
    
with gr.Blocks() as demo:
    gr.Markdown(
        """
        <div style="display: flex;align-items: center;">
            <img alt="Opulentyn Logo" src="https://github.com/user-attachments/assets/5fc78032-bff2-4f7e-a174-7d64b22f506d" width="350"/>
            <div>
                <h1>ID Document Recognition</h1>
                <p>We offer <b>on-premises</b> OCR and liveness check solutions available with a <b>perpetual license</b>.</p>
            </div>
        </div>
        
        ## 🤝 Talk to us
        
        <div style="display: flex; align-items: center;">
            <a href="https://opulentyn.com" target="_blank">
                <img src="https://img.shields.io/badge/Website-https%3A%2F%2Fopulentyn.com-blue?style=flat&logo=google-chrome&logoColor=white" alt="Website">
            </a>
            &nbsp;&nbsp;&nbsp;&nbsp;
            <a href="mailto:support@opulentyn.com">
                <img src="https://img.shields.io/badge/Email-support%40opulentyn.com-blue?style=flat&logo=gmail&logoColor=white" alt="Email">
            </a>
            &nbsp;&nbsp;&nbsp;&nbsp;
            <a href="https://join.slack.com/t/opulentyn/shared_invite/zt-2s230jtbq-dWBs8XUZcrYim~nUqiimSA" target="_blank">
                <img src="https://img.shields.io/badge/Slack-support--sdk-blueviolet?style=flat&logo=slack&logoColor=white" alt="Slack">
            </a>
        </div>
        """
    )

    with gr.Row():
        with gr.Column():
            image_input1 = gr.Image(type='pil', label='Front')
            image_input2 = gr.Image(type='pil', label='Back')
            gr.Examples(['examples/1.jpg', 'examples/2.jpg', 'examples/3.jpg'], 
                        inputs=image_input1)
            process_button = gr.Button("Process")
        with gr.Column():
            with gr.TabItem("Output"):
                html_output = gr.HTML()
            with gr.TabItem("JSON"):
                json_output = gr.JSON()
    
    process_button.click(process_image, inputs=[image_input1, image_input2], outputs=[html_output, json_output])

    gr.HTML('<a href="https://visitorbadge.io/status?path=https%3A%2F%2Fhuggingface.co%2Fspaces%2Fopulentyn%2FIDDocRecog"><img src="https://api.visitorbadge.io/api/combined?path=https%3A%2F%2Fhuggingface.co%2Fspaces%2Fopulentyn%2FIDDocRecog&countColor=%23263759" /></a>')

demo.launch(server_name="0.0.0.0", server_port=7860)