import os
import textwrap
import time

import gradio as gr
from PIL import Image

from check import DEFAULT_MODEL, predict, MODELS


def _predict_fn(image: Image.Image, model_name: str = DEFAULT_MODEL, max_batch_size: int = 8):
    start_time = time.time()
    result = predict(image, model_name, max_batch_size)
    duration = time.time() - start_time
    info = f'Time cost: **{duration:.3f}s**'
    return result, info


if __name__ == '__main__':
    with gr.Blocks() as demo:
        with gr.Row():
            gr_info = gr.Markdown(textwrap.dedent("""
                 Quickly check if the image is **glazed or misted** (we call it shat💩).

                 And then you can just remove these shit with
                 [mf666/mist-fucker](https://huggingface.co/spaces/mf666/mist-fucker),
                 without fucking the normal images (no detail losses).
            """).strip())

        with gr.Row():
            with gr.Column():
                gr_input = gr.Image(label='Image To Check', image_mode='RGB', type='pil')
                gr_models = gr.Dropdown(choices=MODELS, value=DEFAULT_MODEL, label='Models')
                gr_max_batch_size = gr.Slider(minimum=1, maximum=16, value=8, step=1, label='Max Batch Size')
                gr_submit = gr.Button(value='Check The Shit', variant='primary')

            with gr.Column():
                gr_label = gr.Label(label='Check Result')
                gr_time_cost = gr.Markdown(label='Information')

            gr_submit.click(
                _predict_fn,
                inputs=[gr_input, gr_models, gr_max_batch_size],
                outputs=[gr_label, gr_time_cost],
            )

    demo.queue(os.cpu_count()).launch()