from typing import Annotated
from typing_extensions import TypedDict
from langgraph.graph import StateGraph, START, END
from langgraph.graph.message import add_messages
from langchain.chat_models import init_chat_model
import os

class State(TypedDict):
    messages: Annotated[list, add_messages]

graph_builder = StateGraph(State)

key = os.getenv("OPENAI_API_KEY")

llm = init_chat_model("openai:gpt-4o-mini")
tool = {"type": "web_search_preview"}
llm_with_tools = llm.bind_tools(
    [
        {
            "type": "mcp",
            "server_label": "clickup_data",
            "server_url": "https://roniorque-fastapi-test-server.hf.space/mcp",
            "require_approval": "never",
        }
    ]
    + [tool]
    
)

def chatbot(state: State):
    return {"messages": [llm_with_tools.invoke(state["messages"])]}

graph_builder.add_node("chatbot", chatbot)
graph_builder.add_edge(START, "chatbot")
graph_builder.add_edge("chatbot", END)
graph = graph_builder.compile()

if __name__ == "__main__":
    print("Chatbot started! Type 'quit' to exit.")
    
    while True:
        user_input = input("\nYou: ")
        
        if user_input.lower() in ['quit', 'exit', 'bye']:
            print("Goodbye!")
            break
            
        # Create initial state with user message
        initial_state = {
            "messages": [{"role": "user", "content": user_input}]
        }
        
        # Run the graph
        result = graph.invoke(initial_state)
        
        # Get the bot's response
        bot_response = result["messages"][-1].content
        print(f"Bot: {bot_response}")