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, 
                 analyze_excel_file)
import os
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, HuggingFaceEmbeddings
from langchain_google_genai import ChatGoogleGenerativeAI
from langchain_groq import ChatGroq
from langchain_community.vectorstores import SupabaseVectorStore
from langchain.tools.retriever import create_retriever_tool
from supabase.client import Client, create_client

HUGGINGFACEHUB_API_TOKEN = getenv("HUGGINGFACEHUB_API_TOKEN")
SUPABASE_URL = os.environ.get("SUPABASE_URL")
SUPABASE_SERVICE_ROLE_KEY = os.environ.get("SUPABASE_SERVICE_ROLE_KEY")


# 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)

# # build a retriever
# embeddings = HuggingFaceEmbeddings(
#     model_name="sentence-transformers/all-mpnet-base-v2"
# )  #  dim=768
# supabase: Client = create_client(
#     SUPABASE_URL, SUPABASE_SERVICE_ROLE_KEY
# )
# vector_store = SupabaseVectorStore(
#     client=supabase,
#     embedding=embeddings,
#     table_name="documents2",
#     query_name="match_documents_2",
# )
# create_retriever_tool = create_retriever_tool(
#     retriever=vector_store.as_retriever(),
#     name="Question Search",
#     description="A tool to retrieve similar questions from a vector store.",
# )

# 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, 
         analyze_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([sys_msg] + state["messages"])]}

    # def retriever(state: MessagesState):
    #     """Retriever node"""
    #     similar_question = vector_store.similarity_search(state["messages"][0].content)

    #     if similar_question:  # Check if the list is not empty
    #         example_msg = HumanMessage(
    #             content=f"Here I provide a similar question and answer for reference: \n\n{similar_question[0].page_content}",
    #         )
    #         return {"messages": [sys_msg] + state["messages"] + [example_msg]}
    #     else:
    #         # Handle the case when no similar questions are found
    #         return {"messages": [sys_msg] + state["messages"]}

    
    # Build graph / nodes
    builder = StateGraph(MessagesState)
    #builder.add_node("retriever", retriever) # Retriever
    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_edge(START, "assistant")
    builder.add_conditional_edges("assistant", tools_condition)
    builder.add_edge("tools", "assistant")
    
    graph = builder.compile()

    return graph