import gradio as gr
from langchain_community.agent_toolkits.load_tools import load_tools
from langchain.agents import initialize_agent, AgentType
from langchain_openai import ChatOpenAI
import warnings

warnings.filterwarnings("ignore", message=".*TqdmWarning.*")


# ----------------------
# Chatbot function
# ----------------------
def chatbot(api_key: str, query: str) -> str:
    """Handles user query and returns agent response using the provided API key."""
    if not api_key.strip():
        return "⚠️ Please enter a valid OpenAI API key."

    try:
        # Initialize model with user-provided key
        llm_model = "gpt-3.5-turbo"
        llm = ChatOpenAI(temperature=0, model=llm_model, openai_api_key=api_key)

        # Load tools
        tools = load_tools(["llm-math", "wikipedia"], llm=llm)

        # Initialize agent
        agent = initialize_agent(
            tools,
            llm,
            agent=AgentType.CHAT_ZERO_SHOT_REACT_DESCRIPTION,
            handle_parsing_errors=True,
            verbose=True
        )

        # Run the query
        response = agent.run(query)
        return response

    except Exception as e:
        return f"❌ Error: {str(e)}"


# ----------------------
# Gradio Interface
# ----------------------
demo = gr.Interface(
    fn=chatbot,
    inputs=[
        gr.Textbox(label="🔑 OpenAI API Key", type="password", placeholder="Paste your API key here..."),
        gr.Textbox(label="Your Question", placeholder="Ask me anything...")
    ],
    outputs=gr.Textbox(label="Response"),
    title="LangChain AI Chatbot",
    description="Enter your OpenAI API key and ask me anything! Powered by OpenAI and LangChain.",
    theme="compact"
)

# ----------------------
# Launch
# ----------------------
if __name__ == "__main__":
    demo.launch(share=True)