#!/usr/bin/env python

from __future__ import annotations

import gradio as gr
import PIL.Image
from genTag import genTag

def predict(image: PIL.Image.Image, score_threshold: float):
    result_threshold = genTag(image, score_threshold)
    result_html = ''
    for label, prob in result_threshold.items():
        result_html = result_html + '<p class="m5dd_list use"><span>' + str(label) + '</span><span>' + str(round(prob, 3)) + '</span></p>'
    result_html = '<div>' + str(result_html) + '</div>'
    result_text = ', '.join(result_threshold.keys())
    result_text = '<div id="m5dd_result">' + str(result_text) + '</div>'
    return result_html, result_text

js = """
async () => {
    document.addEventListener('click', function(event) {
        let tagItem = event.target.closest('.m5dd_list')
        let resultArea = event.target.closest('#m5dd_result')
        if (tagItem){
            if (tagItem.classList.contains('use')){
                tagItem.classList.remove('use')
            }else{
                tagItem.classList.add('use')
            }
            document.getElementById('m5dd_result').innerText =
                Array.from(document.querySelectorAll('.m5dd_list.use>span:nth-child(1)'))
                .map(v=>v.innerText)
                .join(', ')
        }else if (resultArea){
            const selection = window.getSelection()
            selection.removeAllRanges()
            const range = document.createRange()
            range.selectNodeContents(resultArea)
            selection.addRange(range)
        }else{
            return
        }
    })
}
"""

with gr.Blocks(css="style.css") as demo:
    with gr.Row():
        with gr.Column(scale=1):
            image = gr.Image(label='Input', type='pil')
            score_threshold = gr.Slider(label='Score threshold',
                                        minimum=0,
                                        maximum=1,
                                        step=0.05,
                                        value=0.5)
            run_button = gr.Button('Run')
            result_text = gr.HTML(value="<div></div>")
        with gr.Column(scale=3):
            result_html = gr.HTML(value="<div></div>")

    run_button.click(
        fn=predict,
        inputs=[image, score_threshold],
        outputs=[result_html, result_text],
        api_name='predict',
    )
    demo.load(None,None,None,_js=js)

demo.queue().launch()