import os
from dotenv import load_dotenv
from typing import TypedDict, Annotated

from langgraph.graph import MessagesState, START, StateGraph
from langgraph.graph.message import add_messages
from langgraph.prebuilt import tools_condition, ToolNode
from langchain_core.messages import HumanMessage, SystemMessage, AnyMessage
from langchain_core.tools import tool
from langchain.tools.retriever import create_retriever_tool

from langchain_community.tools.tavily_search import TavilySearchResults
from langchain_community.document_loaders import WikipediaLoader
from langchain_community.document_loaders import ArxivLoader

# from langchain_google_genai import ChatGoogleGenerativeAI
from langchain_groq import ChatGroq

from langchain_huggingface import HuggingFaceEmbeddings

from langchain_community.vectorstores import SupabaseVectorStore
from langchain.schema.document import Document
from supabase import create_client, Client

load_dotenv()

#os.environ["TAVILY_API_KEY"] = os.environ.get("TAVILY_API_KEY")
#os.environ["GOOGLE_API_KEY"] = os.environ.get("GOOGLE_API_KEY")
#os.environ["GROQ_API_KEY"] = os.environ.get("GROQ_API_KEY")

__embeddings = HuggingFaceEmbeddings(
    model_name="sentence-transformers/all-mpnet-base-v2",
    model_kwargs= { 'device': 'cpu' })

# connect to supabase
url: str = os.getenv("SUPABASE_URL")
key: str = os.getenv("SUPABASE_SECRET_KEY")
__supabase: Client = create_client(url, key)

# build retriever
vector_store = SupabaseVectorStore(
    client=__supabase,
    embedding=__embeddings,
    table_name="documents",
    query_name="match_documents",
)
question_retrieval_tool = create_retriever_tool(
    vector_store.as_retriever(),
    name="Question retriever",
    description="Find similar questions in the vector database for the given question."
)

# load prompt message from txt file and convert to System Message
with open('prompt.txt', 'r', encoding='utf-8') as f:
    sys_prompt = f.read()

__sys_msg = SystemMessage(content=sys_prompt)

@tool
def add(a: int, b: int) -> int:
    """Add two numbers.
    Args:
        a: first int
        b: second int
    """
    return a + b

@tool
def subtract(a: int, b: int) -> int:
    """Subtract two numbers.
    Args:
        a: first int
        b: second int
    """
    return a - b

@tool
def multiply(a: int, b: int) -> int:
    """Multiply two numbers.
    Args:
        a: first int
        b: second int
    """
    return a * b

@tool
def power(a: int, b: int) -> int:
    """Power up first number by second number.
    Args:
        a: first int
        b: second int
    """
    return a ** b

@tool
def divide(a: int, b: int) -> int:
    """Divide first number by second number.
    Args:
        a: first int
        b: second int
    """
    try:
        return a / b
    except ZeroDivisionError:
        return None
    
@tool
def modulus(a: int, b: int) -> int:
    """Get remainder of first number divided by second number.
    Args:
        a: first int
        b: second int
    """
    return a % b

@tool 
def wiki_search(query: str) -> str:
    """Search Wikipedia for a query and return maximum 2 results.
    
    Args:
        query: The search query.
    """
    search_docs = WikipediaLoader(query=query, load_max_docs=2).load()
    formatted_search_docs = "\n\n---\n\n".join([
        f'<Document source="{doc.metadata["source"]}" page="{doc.metadata.get("page", "")}"/>\n\t{doc.page_content}\n<Document>'
        for doc in search_docs
    ])
    return { "wiki_results": formatted_search_docs }

@tool 
def web_search(query: str) -> str:
    """Search Tavily for a query and return maximum 3 results.
    
    Args:
        query: The search query.
    """
    search_docs = TavilySearchResults(max_results=3).invoke(input=query)
    formatted_search_docs = "\n\n---\n\n".join([
        f'<Document source="{doc["url"]}"/>\n\t{doc["content"]}\n<Document>'
        for doc in search_docs
    ])
    return { "web_results": formatted_search_docs }

@tool 
def arxiv_search(query: str) -> str:
    """Search Arxiv for a query and return maximum 3 result.
    
    Args:
        query: The search query.
    """
    search_docs = ArxivLoader(query=query, load_max_docs=3).load()
    formatted_search_docs = "\n\n---\n\n".join([
        f'<Document source="{doc.metadata["source"]}" page="{doc.metadata.get("page", "")}"/>\n\t{doc.page_content[:1000]}\n<Document>'
        for doc in search_docs
    ])
    return { "arxiv_results": formatted_search_docs }


# list of tools
tools = [
    add,
    subtract,
    multiply,
    power,
    divide,
    modulus,
    wiki_search,
    web_search,
    arxiv_search
]

# Generate the AgentState and Agent graph
class AgentState(TypedDict):
    messages: Annotated[list[AnyMessage], add_messages]


def build_graph():
    # llm = ChatGoogleGenerativeAI(model="gemini-2.0-flash", temperature=0)
    llm = ChatGroq(model="qwen-qwq-32b", temperature=0)
    llm_with_tools = llm.bind_tools(tools)

    # Node 
    def assistant(state: AgentState):
        """Assistant node"""
        return { "messages": [llm_with_tools.invoke(state['messages'])] }
    
    def retriever(state: AgentState):
        similar_question = vector_store.similarity_search(state['messages'][0].content)
        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] }

    builder = StateGraph(AgentState)

    # Define nodes: these do the work
    builder.add_node("assistant", assistant)
    builder.add_node("retriever", retriever)
    builder.add_node("tools", ToolNode(tools))
    builder.add_edge(START, "retriever")
    builder.add_conditional_edges(
        "assistant",
        tools_condition
    )
    builder.add_edge("tools", "assistant")
    builder.add_edge("retriever", "assistant")

    # Compile graph
    return builder.compile()

# Test
if __name__ == "__main__":
    # question = "When was a picture of St. Thomas Aquinas first added to the Wikipedia page on the Principle of double effect?"
    question = "How many studio albums were published by Mercedes Sosa between 2000 and 2009 (included)? You can use the latest 2022 version of english wikipedia."
    graph = build_graph()
    messages = [HumanMessage(content=question)]
    messages = graph.invoke({ "messages": messages })
    for m in messages["messages"]:
        m.pretty_print()