import gradio as gr
import time
import pandas

with gr.Blocks() as demo:
    with gr.Row():
        chatbot = gr.Chatbot()
        with gr.Column(variant="panel"):
            with gr.Row():
                line_count = gr.Number(50, label="Line Count")
                content = gr.Radio(["Text", "Markdown", "Code"], value="Text", label="Content Type to Generate")
            generate_btn = gr.Button("Generate")

    durations_plot = gr.LinePlot(x="line_count", y="duration_ms", height=400, width=600)    
    durations = gr.State([])

    def generate(history, line_count, content, durations):
        start_time = time.time()
        line_count = int(line_count)
        for i in range(line_count):
            if i != 0:
                durations.append((time.time() - start_time) * 1000)
                start_time = time.time()
            if content == "Text":
                yield history + [[None, "lorem ipsum\n" * i]]
            elif content == "Markdown":
                yield history + [[None, "- markdown!\n" * i]]
            elif content == "Code":
                yield history + [[None, "```python\n" + "foo = bar()\n" * i + "```"]]

    def plot(durations):
        df = pandas.DataFrame({"duration_ms": durations, "line_count": list(range(1, len(durations) + 1))})
        return df

    generate_btn.click(generate, [chatbot, line_count, content, durations], chatbot).then(
        plot, durations, durations_plot)

    
if __name__ == "__main__":
    demo.queue().launch()