import gradio as gr
import time

with gr.Blocks() as demo:
    with gr.Row():
        messages_per_second = gr.Number(100, label="Messages per second")
        number_of_messages = gr.Number(200, label="Number of messages")

    chatbot = gr.Chatbot()
    start = gr.Button("Start")

    def bot(messages_per_second, number_of_messages):
        history = [[f"Streaming {number_of_messages} msgs at {messages_per_second} msg/s ", ""]]
        start_time = time.time()
        for num in range(number_of_messages):
            history[-1][1] += " " + str(num)
            time.sleep(1 / messages_per_second)
            yield history, None
        end_time = time.time()
        yield history, end_time - start_time
        
    
    with gr.Row():
        expected_duration = gr.Number(label="Expected Duration (w/o Overhead)")
        backend_duration = gr.Number(label="Duration (for Backend)")
        frontend_duration = gr.Number(label="Duration (for Frontend)")

    start.click(fn=None, js="""() => {
        window.start_time = performance.now();
    }""")
    start.click(lambda messages_per_second, number_of_messages: number_of_messages / messages_per_second, [messages_per_second, number_of_messages], [expected_duration])
    start.click(bot, [messages_per_second, number_of_messages], [chatbot, backend_duration]).then(
        fn=None,
        outputs=frontend_duration,
        js="""() => {
        let now = performance.now();
        return (now - window.start_time) / 1000;                                                                                    
    }""")

    
demo.launch()