import os
from dotenv import load_dotenv
from datetime import datetime, timezone
from langchain_openai import ChatOpenAI
from langchain_core.tools import Tool
from langchain_community.utilities import SerpAPIWrapper
from langchain_core.messages import HumanMessage
from langgraph.prebuilt import create_react_agent
import gradio as gr

# Load environment variables
load_dotenv(override=True)

MODEL = "gpt-4o"

# Optional: real-time search via SerpAPI (set SERPAPI_API_KEY in .env to enable)
if os.environ.get("SERPAPI_API_KEY"):
    search = SerpAPIWrapper(params={"num": "10", "hl": "en", "gl": "in"})
    tools = [
        Tool(
            name="google_search",
            func=search.run,
            description="Use this tool ONLY to fetch real-time data or current events. Always add today's date to ensure accuracy. Never rely on prior knowledge or make assumptions about ongoing events."
        )
    ]
else:
    tools = []  # App runs with GPT only; add SERPAPI_API_KEY to .env for search

# Initialize the LLM
llm = ChatOpenAI(temperature=0, model=MODEL)

# Create ReAct agent only when search is available (create_react_agent requires at least one tool)
agent = create_react_agent(llm, tools) if tools else None

# Function to inject current date into query if missing
def add_current_date_to_query(query):
    today = datetime.now(timezone.utc).strftime("%B %d, %Y %H:%M")
    if today.lower() not in query.lower():
        query = f"{query} ({today})"
    return query

# Chatbot function with exception handling and enforced date
def chatbot(user_input):
    if user_input.lower().strip() in ["hi", "hello", "hey"]:
        return "Hello! How can I assist you today?"
    try:
        enhanced_input = add_current_date_to_query(user_input)
        if agent is not None:
            result = agent.invoke({"messages": [HumanMessage(content=enhanced_input)]})
            messages = result.get("messages", [])
            output = messages[-1].content if messages else "No response generated."
        else:
            # No search tool: call LLM directly (add SERPAPI_API_KEY to .env for real-time search)
            response = llm.invoke([HumanMessage(content=enhanced_input)])
            output = response.content if hasattr(response, "content") else str(response)
    except Exception as e:
        output = f"Unexpected Error: {e}"
    return output

# Gradio UI with chat history
with gr.Blocks(theme=gr.themes.Soft()) as demo:
    gr.HTML("""
    <style>
        body, .gr-block, .gr-button, .gr-textbox, .gr-chatbot {
            font-family: 'Times New Roman', Times, serif !important;
            font-size: 16px;
        }
    </style>
    """)
    # Header and description
    gr.Markdown(
        """
        # Smart AI Assistant  
        Powered by **GPT + Real-time Search**.
        * By - **Manoj Anuragi**
        """,
        elem_id="header"
    )

    with gr.Row():
        chatbot_ui = gr.Chatbot(
            type="messages",
            autoscroll=True,
            height=450,
            show_label=False
        )
    
    with gr.Row():
        user_input = gr.Textbox(
            placeholder="Ask me anything...",
            show_label=False,
            scale=10,

            autofocus=True
        )
    submit = gr.Button(" ➡️ Ask",  scale=2)

    def respond(user_message, chat_history):
        chat_history = chat_history or []
        chat_history.append({"role": "user", "content": user_message})
        response = chatbot(user_message)
        chat_history.append({"role": "assistant", "content": response})
        return chat_history, ""

    user_input.submit(respond, inputs=[user_input, chatbot_ui], outputs=[chatbot_ui, user_input])
    submit.click(respond, inputs=[user_input, chatbot_ui], outputs=[chatbot_ui, user_input])

# Launch app
demo.launch(inline=False)