import gradio as gr
import time

sleep_time = 0.5

def user_message(message, history):
    """Add user message to chat history"""
    return "", history + [{"role": "user", "content": message}]

def simulate_thinking_chat(history):
    """Generate bot responses"""
    start_time = time.time()
    
    # Initialize thinking response
    thinking_header = "_Thinking_ step-by-step"
    history.append({"role": "assistant", "content": f"**{thinking_header}**\n\n"})
    yield history

    # Generate thoughts progressively
    thoughts = [
        "First, I need to understand the core aspects of the query...",
        "Now, considering the broader context and implications...",
        "Analyzing potential approaches to formulate a comprehensive answer...",
        "Finally, structuring the response for clarity and completeness..."
    ]

    accumulated_thoughts = ""
    for thought in thoughts:
        time.sleep(sleep_time)
        accumulated_thoughts += f"- {thought}\n\n"
        history[-1]["content"] = f"**{thinking_header}**\n\n{accumulated_thoughts.strip()}"
        yield history

    # Add completion note with duration
    duration = time.time() - start_time
    history[-1]["content"] = f"**{thinking_header}** _(completed in {duration:.2f}s)_\n\n{accumulated_thoughts.strip()}"
    yield history

    # Add final response with a delay to showcase loading animation
    time.sleep(5.0)
    history.append({
        "role": "assistant", 
        "content": "Based on my thoughts and analysis above, my response is: This dummy repro shows how thoughts of a thinking LLM can be progressively shown before providing its final answer."
    })
    yield history

with gr.Blocks(title="Thinking LLM Chat Interface 🤔") as demo:

    chatbot = gr.Chatbot(type="messages")
    msg = gr.Textbox(label="Message", placeholder="Type your message here...")
    clear = gr.ClearButton([msg, chatbot])

    msg.submit(user_message, [msg, chatbot], [msg, chatbot]).then(
        simulate_thinking_chat, chatbot, chatbot)

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