import datetime
import pytz
import re
import gradio as gr

def get_time(timezone_str: str) -> str:
    try:
        tz = pytz.timezone(timezone_str)
    except pytz.UnknownTimeZoneError:
        return f"Sorry, '{timezone_str}' is not a recognized timezone."
    now = datetime.datetime.now(tz)
    return now.strftime("%Y-%m-%d %H:%M:%S")

def agent_response(user_input: str) -> str:
    # Updated regex to capture timezone after 'time' and 'in'
    match = re.search(r"time\s+(?:is\s+it\s+)?in\s+([\w/_]+)", user_input, re.IGNORECASE)
    if match:
        timezone_str = match.group(1)
        time_str = get_time(timezone_str)
        return f"The current time in {timezone_str} is {time_str}"
    else:
        return ("Sorry, I can only answer questions about time in a specific timezone, "
                "e.g. 'What time is it in America/New_York?'")

def main():
    with gr.Blocks() as demo:
        gr.Markdown("## Ask me about the current time in a timezone!")
        chatbot = gr.Chatbot()
        txt = gr.Textbox(placeholder="Enter your question here...")

        def respond(user_message, chat_history):
            response = agent_response(user_message)
            chat_history = chat_history + [(user_message, response)]
            return chat_history, ""

        txt.submit(respond, inputs=[txt, chatbot], outputs=[chatbot, txt])
        txt.submit(lambda: "", None, txt)  # Clear input box after submit

    demo.launch()

if __name__ == "__main__":
    main()