# LangGraph Agen
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_groq import ChatGroq
from langchain_huggingface import ChatHuggingFace, HuggingFaceEndpoint, HuggingFaceEmbeddings
from langchain_community.tools.tavily_search.tool import TavilySearchResults
from langchain_community.document_loaders import WikipediaLoader
from langchain_community.document_loaders import ArxivLoader
from langchain_community.vectorstores import SupabaseVectorStore
from langchain_core.messages import SystemMessage, HumanMessage
from langchain_core.tools import tool
from langchain.tools.retriever import create_retriever_tool
from supabase.client import Client, create_client
from langchain_community.utilities import GoogleSerperAPIWrapper
from langchain.schema import Document
from youtubeAnalyseTool import YoutubeSearchTool
from langchain_core.messages import AIMessage

load_dotenv()

youtube_api_key = os.getenv("YOUTUBE_API_KEY")

# Tool Decleration

@tool
def calculate_sum_of_two_integers(number1: int, number2: int) -> int:
    """
    Calculate sum of two integers
    args:
        number1: first integer 
        number2: second integer
    """
    return number1 + number2

@tool
def multiply_two_integers(number1: int, number2: int) -> int:
    """
    Multiply two integers
    args: 
        number1: first integer
        number2: second integer
    """
    return number1 * number2

@tool 
def divide_two_integers(number1: int, number2: int) -> float:
    """
    Divide two integers
    args:
        number1: first integer
        number2: second integer
    """
    result = float(number1) / float(number2)
    return round(result, 2)

@tool
def subtract_two_integers(number1: int, number2: int) -> int:
    """
    Calculate difference between two integers
    args:
        number1: first integer
        number2: second integer
    """
    return number1 - number2

@tool 
def search_wikipedia(searchTerm: str) -> str:
    """
    Search wikipedia for a query and return maximum 3 results 
    args:
        searchTerm: the search term to query
    """
    query = WikipediaLoader(query = searchTerm, load_max_docs=3).load()
    format_response = "\n\n---\n\n".join([
            f'<Document source="{doc.metadata["source"]}" page="{doc.metadata.get("page","")}"/>\n{doc.page_content}\n</Document>'
            for doc in query
        ])
    return {"Result from wikipedia" : format_response}

@tool
def web_search_tavily(searchTerm: str) -> str:
    """
    Search tavily for a query and return maximum of 3 results
    args:
        searchTerm: the search term to query
    """
    query = TavilySearchResults(max_results=3).invoke(query=searchTerm)
    format_response = "\n\n---\n\n".join([
        f'<Document source="{doc.metadata["source"]}" page="{doc.metadata.get("page", "")}"/>\n{doc.page_content}\n</Document>'
        for doc in query
    ])
    return {"Result from Tavily": format_response}

@tool 
def get_news_from_google(searchTerm: str) -> str:
    """
    Search for news on google and return a maximum of top 3 relevant results
    args:
        searchTerm: the search term to query
    """
    wrapper = GoogleSerperAPIWrapper(type="news")
    result = wrapper.results(searchTerm)
    top_3_results = result.get("news", [])[:3]

    if not top_3_results:
        return f"No relevant news found for the query : {searchTerm}"

    docs = []
    for news in top_3_results:
        content = content = f"{news.get('title','')}\n{news.get('snippet','')}\nURL:{news.get('link','')}"
        metadata = {"source":news.get("link",""), "page":""}
        docs.append(Document(page_content=content,metadata=metadata))

    format_response = "\n\n---\n\n".join([
        f'<Document source="{doc.metadata["source"]}" page="{doc.metadata.get("page", "")}">\n{doc.page_content}\n</Document>'
        for doc in docs
    ])
    return {"News from google": format_response}

youtube_tool = YoutubeSearchTool(youtube_api_key=youtube_api_key)

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

# set as system message 
sys_message = SystemMessage(content=system_prompt)

# define tools
tools = [
    calculate_sum_of_two_integers,
    multiply_two_integers,
    divide_two_integers,
    subtract_two_integers,
    search_wikipedia,
    web_search_tavily,
    get_news_from_google,
    youtube_tool,
]

# Build graph function
def build_graph(provider: str = "google"):
    """Build the graph"""

    if provider == "google":
        llm = ChatHuggingFace(
            llm=HuggingFaceEndpoint(
                repo_id="meta-llama/Llama-2-7b-chat-hf",
                temperature=0,
            ),
        )
    elif provider == "groq":
        llm = ChatHuggingFace(
            llm=HuggingFaceEndpoint(
                repo_id="meta-llama/Llama-2-7b-chat-hf",
                temperature=0,
            ),
        )
    elif provider == "huggingface":
        llm = ChatHuggingFace(
            llm=HuggingFaceEndpoint(
                repo_id="meta-llama/Llama-2-7b-chat-hf",
                temperature=0,
            ),
        )
    else:
        raise ValueError("Invalid provider. Choose 'google', 'groq' or 'huggingface'.")

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

    def assistant(state: MessagesState):
        """Assistant node"""
        return {"messages": [llm_with_tools.invoke(state["messages"])]}

    # Build the graph
    builder = StateGraph(MessagesState)
    builder.add_node("assistant", assistant)
    builder.add_node("tool_use", ToolNode(tools=tools))

    builder.add_conditional_edges(
        "assistant",
        tools_condition,
        path_map={"tool_use": "tool_use"}
    )
    builder.add_edge("tool_use", "assistant")
    builder.set_entry_point("assistant")
    builder.set_finish_point("assistant")

    return builder.compile()