import gradio as gr

from .defaults import defaults_vanilla_single, defaults_bt_spam, defaults_bt_single, defaults_vanilla_spam
from .defaults import ADDRESS_VANILLA, ADDRESS_BETTERTRANSFORMER
from .backend import send_single, send_spam, get_message_single, get_message_spam

with gr.Blocks() as demo:
    gr.Markdown("""
    Let's try out TorchServe + BetterTransformer! This is some longer description This is some longer description This is some longer description")
    
    ## Inference using...
    """
    )
    
    with gr.Row():
        with gr.Column(scale=50):
            gr.Markdown("### Vanilla Transformers + TorchServe")

            address_input_vanilla = gr.Textbox(
                max_lines=1,
                label="ip vanilla",
                value=ADDRESS_VANILLA,
                visible=False
            )
            
            input_model_vanilla = gr.Textbox(
                max_lines=1,
                label="Text",
                value="Expectations were low, enjoyment was high",
            )

            btn_single_vanilla = gr.Button("Send single text request")
            output_single_vanilla = gr.Markdown(label="Output single vanilla", value=get_message_single(**defaults_vanilla_single))
            
            btn_spam_vanilla = gr.Button("Spam text requests (from sst2 validation set)") 
            output_spam_vanilla = gr.Markdown(label="Output spam vanilla", value=get_message_spam(**defaults_vanilla_spam))
            
            btn_single_vanilla.click(
                fn=send_single,
                inputs=[input_model_vanilla, address_input_vanilla],
                outputs=output_single_vanilla,
            )
            btn_spam_vanilla.click(
                fn=send_spam,
                inputs=[address_input_vanilla],
                outputs=output_spam_vanilla,
            )

        with gr.Column(scale=50):
            gr.Markdown("### BetterTransformer + TorchServe")

            address_input_bettertransformer = gr.Textbox(
                max_lines=1,
                label="ip bettertransformer",
                value=ADDRESS_BETTERTRANSFORMER,
                visible=False
            )
            
            input_model_bettertransformer = gr.Textbox(
                max_lines=1,
                label="Text",
                value="Expectations were low, enjoyment was high",
            )

            btn_single_bt = gr.Button("Send single text request")
            output_single_bt = gr.Markdown(label="Output single bt", value=get_message_single(**defaults_bt_single))
            
            btn_spam_bt = gr.Button("Spam text requests (from sst2 validation set)")
            output_spam_bt = gr.Markdown(label="Output spam bt", value=get_message_spam(**defaults_bt_spam))
                        
            btn_single_bt.click(
                fn=send_single,
                inputs=[input_model_bettertransformer, address_input_bettertransformer],
                outputs=output_single_bt,
            )
            
            btn_spam_bt.click(
                fn=send_spam,
                inputs=[address_input_bettertransformer],
                outputs=output_spam_bt,
            )

demo.queue(concurrency_count=1)
demo.launch()