import os
from dotenv import load_dotenv
from langgraph.graph import START, StateGraph, MessagesState
from langgraph.prebuilt import tools_condition
from langgraph.prebuilt import ToolNode
from langchain_google_genai import ChatGoogleGenerativeAI
from langchain_huggingface import ChatHuggingFace, HuggingFaceEndpoint, HuggingFaceEmbeddings
from langchain_community.tools.tavily_search import TavilySearchResults
from langchain_community.document_loaders import WikipediaLoader, ArxivLoader
from langchain_community.vectorstores import SupabaseVectorStore
from langchain_core.messages import SystemMessage, AIMessage, HumanMessage
from langchain_core.tools import tool
from langchain.tools.retriever import create_retriever_tool
from supabase.client import Client, create_client

# Load environment variables
load_dotenv()

# --- Math Tools ---
@tool
def multiply(a: int, b: int) -> int:
    """Multiply two integers."""
    return a * b

@tool
def add(a: int, b: int) -> int:
    """Add two integers."""
    return a + b

@tool
def subtract(a: int, b: int) -> int:
    """Subtract b from a."""
    return a - b

@tool
def divide(a: int, b: int) -> float:
    """Divide a by b, error on zero."""
    if b == 0:
        raise ValueError("Cannot divide by zero.")
    return a / b

@tool
def modulus(a: int, b: int) -> int:
    """Compute a mod b."""
    return a % b

# --- Browser Tools ---
@tool
def wiki_search(query: str) -> dict:
    """Search Wikipedia and return up to 2 documents."""
    docs = WikipediaLoader(query=query, load_max_docs=2).load()
    results = [f"<Document source=\"{d.metadata['source']}\" page=\"{d.metadata.get('page','')}\"/>\n{d.page_content}" for d in docs]
    return {"wiki_results": "\n---\n".join(results)}

@tool
def web_search(query: str) -> dict:
    """Search Tavily and return up to 3 results."""
    docs = TavilySearchResults(max_results=3).invoke(query=query)
    results = [f"<Document source=\"{d.metadata['source']}\" page=\"{d.metadata.get('page','')}\"/>\n{d.page_content}" for d in docs]
    return {"web_results": "\n---\n".join(results)}

@tool
def arxiv_search(query: str) -> dict:
    """Search Arxiv and return up to 3 docs."""
    docs = ArxivLoader(query=query, load_max_docs=3).load()
    results = [f"<Document source=\"{d.metadata['source']}\" page=\"{d.metadata.get('page','')}\"/>\n{d.page_content[:1000]}" for d in docs]
    return {"arxiv_results": "\n---\n".join(results)}

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

# --- System message ---
sys_msg = SystemMessage(content=system_prompt)

# --- Retriever Tool ---
embeddings = HuggingFaceEmbeddings(model_name="sentence-transformers/all-mpnet-base-v2")
supabase = create_client(os.getenv("SUPABASE_URL"), os.getenv("SUPABASE_SERVICE_KEY"))

vector_store = SupabaseVectorStore(
    client=supabase,
    embedding=embeddings,                                  table_name="documents",
    query_name="match_documents_langchain")

retriever_tool = create_retriever_tool(
    retriever=vector_store.as_retriever(

        search_type="similarity",

        search_kwargs={"k": 5}
    ),
    name="Question Search",
    description="A tool to retrieve similar questions from the vector store."
)

tools = [
    multiply,
    add,
    subtract,
    divide,
    modulus,
    wiki_search,
    web_search,
    arxiv_search,
]

# --- Graph Builder ---
def build_graph(provider: str = "huggingface"):
    llm = ChatHuggingFace(

        llm=HuggingFaceEndpoint(

            repo_id = "Qwen/Qwen2.5-Coder-32B-Instruct"

        ),

    )

    # Bind tools to LLM
    llm_with_tools = llm.bind_tools(tools)

    # Define no    def assistant(state: MessagesState):
        """Assistant node"""

      
          return {"messages [ [llm_with_tools.invoke(state["messages"])]}se]}


    # Retriever  returns AIMessage    def retriever(state: MessagesState):

        """Retriever node"""

        similar_question = vector_store.similarity_search(state["messages"][0].content)

        print('Similar questions:')

        print(similar_question)

        if len(similar_question) > 0:

            example_msg = HumanMessage(

                content=f"Here I provide a similar question and answer for reference: \n\n{similar_question[0].page_content}",

      ntent}]}

            return {"messages": [sys_msg] + state["messages"] + [example_msg]}

        return {"messages": [sys_msg] + state["m
    # Add nodesessages"]}



    builder = StateGraph(MessagesState)

    builder.add_node("retriever", retriever)

    builder.add_node("assistant", assistant)

    builder.add_node("tools", 

    # Add edgesToolNode(tools))

    builder.add_edge(START, "retriever")

    builder.add_edge("retriever", "assistant")

    builder.add_conditional_edges(

        "assistant",

        tools_condition,

    )

    builder.add_edge("tools", "assistant")ever")

    # Compile graph
    return builder.compile()