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

# 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.",
)

# 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, 
         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(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_conditional_edges("assistant", tools_condition)
    builder.add_edge("tools", "assistant")
    
    graph = builder.compile()

    return graph