from typing import TypedDict, Annotated
from tool import (add, 
                 substract, 
                 multiply, 
                 divide, 
                 DuckDuckGoSearchTool, 
                 TavilySearchTool, 
                 WikipediaSearchTool, 
                 ArxivSearchTool, 
                 PubmedSearchTool, 
                 save_and_read_file,
                 download_file_from_url,
                 extract_text_from_image,
                 analyze_csv_file, 
                 analyse_excel_file)
from os import getenv
from langgraph.graph.message import add_messages
from langchain_core.messages import AnyMessage, SystemMessage, HumanMessage, AIMessage
from langgraph.graph import StateGraph, START, END, MessagesState
from langgraph.prebuilt import ToolNode, tools_condition
from langchain_huggingface import HuggingFaceEndpoint, ChatHuggingFace
from langchain_google_genai import ChatGoogleGenerativeAI
from langchain_groq import ChatGroq


HUGGINGFACEHUB_API_TOKEN = getenv("HUGGINGFACEHUB_API_TOKEN")


# load the system prompt from the file
with open("prompt.txt", "r", encoding="utf-8") as f:
    system_prompt = f.read()

# System message
sys_msg = SystemMessage(content=system_prompt)

# Loading the assistant
chat = ChatGoogleGenerativeAI(model="gemini-2.0-flash", temperature=0)

tools = [add, 
         substract, 
         multiply, 
         divide, 
         DuckDuckGoSearchTool, 
         TavilySearchTool, 
         WikipediaSearchTool, 
         ArxivSearchTool, 
         PubmedSearchTool, 
         save_and_read_file,
         download_file_from_url,
         extract_text_from_image,
         analyze_csv_file, 
         analyse_excel_file]

chat_with_tools = chat.bind_tools(tools)

def simple_graph():

    ## Defining our nodes
    def assistant(state: MessagesState):
        """Assistant node"""
        return {"messages": [chat_with_tools.invoke(state["messages"])]}

    def retriever(state: MessagesState):
        """Retriever node"""
        # I don't want to use a Retriever / Using similar questions.
        return {"messages": [sys_msg] + state["messages"]}
    
    # Build graph / nodes
    builder = StateGraph(MessagesState)
    builder.add_node("retriever", retriever) # Assistant
    builder.add_node("assistant", assistant) # Assistant
    builder.add_node("tools", ToolNode(tools)) # Tools
    
    # Logic / edges
    builder.add_edge(START, "retriever")
    builder.add_edge("retriever", "assistant")
    builder.add_conditional_edges("assistant", tools_condition)
    builder.add_edge("tools", "assistant")
    
    graph = builder.compile()

    return graph