import gradio as gr
from llm_handler import sanchaari
import time

def handle_query(query):
    thoughts = []
    final_answer = None
    for result in sanchaari(query):
        if "thought" in result:
            thoughts.append(result["thought"])
            yield query, "\n".join(thoughts), ""
        if "final_answer" in result:
            final_answer = result["final_answer"]
            yield query, "\n".join(thoughts), final_answer

def clear_outputs():
    start_time = time.time()
    result = ("", "", "")
    end_time = time.time()
    print(f"Clear operation took {end_time - start_time:.6f} seconds")
    return result

with gr.Blocks(theme=gr.themes.Soft()) as demo:
    gr.Markdown(
        """
        # Antarjaala Sanchaari: Your Versatile AI Assistant
        
        Meet **Antarjaala Sanchaari**, your intelligent companion for insightful answers and seamless assistance.
        """
    )
    
    with gr.Row():
        with gr.Column(scale=1):
            gr.Markdown(
                """
                ### What Sanchaari Can Do:
                - **Calculate**: Perform arithmetic calculations
                - **Find Information**: Search and summarize from Wikipedia
                - **Tell Time**: Get current time worldwide
                - **Get Stock Prices**: Check current stock prices
                - **Define Words**: Look up word definitions
                - **Provide Latest News**: Retrieve news headlines
                - **Translate Text**: Translate from English
                """
            )
        with gr.Column(scale=1):
            gr.Markdown(
                """
                ### Example Questions:
                1. What is the capital of France?
                2. Calculate 2*3
                3. Define "serendipity"
                4. Translate "Hello, how are you?" to Spanish
                5. What's the current stock price of Apple?
                6. What's the latest news about climate change?
                7. What time is it in Tokyo right now?
                """
            )
    
    with gr.Row():
        with gr.Column(scale=1):
            query_input = gr.Textbox(
                lines=5, 
                placeholder="Enter your query here...",
                label="Your Question"
            )
            with gr.Row():
                submit_btn = gr.Button("Ask Sanchaari", variant="primary")
                clear_btn = gr.Button("Clear", variant="secondary")
            response_output = gr.Textbox(label="Sanchaari's Response", lines=3)
        thought_output = gr.Textbox(label="Thought Process", lines=10, scale=1)
    
    submit_btn.click(handle_query, inputs=query_input, outputs=[query_input, thought_output, response_output])
    clear_btn.click(clear_outputs, outputs=[query_input, thought_output, response_output])

demo.launch(share=True)