from typing import TypedDict, Annotated, Optional
from tools import multiply, sum, substract, divide\
    , wiki_search, web_search, arvix_search, gaia_database_search
    # , read_python, read_excel, transcribe_audio
from llm import chat_llm
from vector_store import get_supabase_client, get_gaia_data_vector_store, get_gaia_data_retriever_tool
from langchain_core.messages import AnyMessage, SystemMessage, HumanMessage
from langgraph.graph.message import add_messages
from langgraph.graph import START, StateGraph
from langgraph.prebuilt import ToolNode, tools_condition

tools = [
    multiply,
    sum,
    substract,
    divide,
    wiki_search,
    web_search,
    arvix_search,
    # read_python,
    # read_excel,
    # transcribe_audio,
    # gaia_database_search
]
llm_with_tools = chat_llm.bind_tools(
    tools
    , parallel_tool_calls=False
)

textual_description_of_tool="""
divide(a: int, b: int) -> float:
    Divide a and b
    Args:
        a: The numerator (int).
        b: The denominator (int).
    Returns:
        The result of the division (float).
multiply(a: float, b: float) -> float:
    multiply a and b
    Args:
        a: Number one.
        b: Number two.
    Returns:
        The result of the multiplication (float).
sum(a: float, b: float) -> float:
    multiply a and b
    Args:
        a: Number one.
        b: Number two.
    Returns:
        The result of the summation (float).
substract(a: float, b: float) -> float:
    multiply a and b
    Args:
        a: Number one.
        b: Number two.
    Returns:
        The result of the substraction (float).
wiki_search(query: str) -> str:
    Search Wikipedia for a query and return maximum 2 results.
    Args:
        query: The Search Query.
    Returns:
        Formatted search result.
web_search(query: str) -> str:
    Search Tavily for a query and return maximum 3 results.
    Args:
        query: The Search Query.
    Returns:
        Formatted search result.
arvix_search(query: str) -> str:
    Search Arxiv for a query and return maximum 3 result.
    Args:
        query: The Search Query.
    Returns:
        Formatted search result.
gaia_database_search(query: str) -> str:
    Search the Supabase vector store for documents relevant to the query.
    Args:
        query: The Search Query.
    Returns:
        Formatted string with the top results
"""
# SYS_MSG = SystemMessage(content=f"""You are a helpful butler named Jarvis that serves Mr. Sarkar.
#     You can search internet, arvix and wikipedia and run computations with provided tools:\n{textual_description_of_tool} \n.
#     Now, I will ask you a question. Always prioritize arvix search over web serach tool and then web search over wiki search whenever possible.
#     - ALWAYS PASS THE HUMAN PROVIDED QUERY TO THE TOOLS, DO NOT SUMMARIZE THE QUESTION BEFORE SENDING TO TOOLS.
#     - For any kind of academic and educational question, use the tool for Arvix search.
#     - Always prioritize internet search tools to get an answer
#     - If answer is not found via internet search tools, use GAIA database search tool to look into provided knowledge
#     Report your thoughts, and finish your answer with the following template:
#     FINAL ANSWER: [YOUR FINAL ANSWER]. 
#     YOUR FINAL ANSWER should be a number OR as few words as possible OR a comma separated list of numbers and/or strings.
#     If you are asked for a number, don't use comma to write your number neither use units such as $ or percent sign unless specified otherwise.
#     If you are asked for a string, don't use articles, neither abbreviations (e.g. for cities), and write the digits in plain text unless specified otherwise.
#     If you are asked for a comma separated list, apply the above rules depending of whether the element to be put in the list is a number or a string.
#     Your answer should only start with "FINAL ANSWER: ", then follows with the answer. 
# """)
SYS_MSG = SystemMessage(content=f"""You are a helpful butler named Jarvis that serves Mr. Sarkar.
    You can search internet, arvix and wikipedia and run computations with binded tools.
    Now, I will ask you a question. Always prioritize arvix search over web serach tool and then web search over wiki search whenever possible.
    - IT IS ABSOLUTELY MANDATORY TO CALL TOOLS WITH USER'S FULL  MESSAGE AS THE `QUERY` ARGUMENT.
    - If answer is not found via internet search tools, use GAIA database search tool to look into provided knowledge
    Report your thoughts, and finish your answer with the following template:
    FINAL ANSWER: [YOUR FINAL ANSWER]. 
    YOUR FINAL ANSWER should be a number OR as few words as possible OR a comma separated list of numbers and/or strings.
    If you are asked for a number, don't use comma to write your number neither use units such as $ or percent sign unless specified otherwise.
    If you are asked for a string, don't use articles, neither abbreviations (e.g. for cities), and write the digits in plain text unless specified otherwise.
    If you are asked for a comma separated list, apply the above rules depending of whether the element to be put in the list is a number or a string.
    Your answer should only start with "FINAL ANSWER: ", then follows with the answer. 
""")


class AgentState(TypedDict):
    # The document provided
    input_file: Optional[str]  # Contains file path (PDF/PNG)
    messages: Annotated[list[AnyMessage], add_messages]


def assistant(state: AgentState):
    # System message

    image=state["input_file"]

    response = llm_with_tools.invoke([SYS_MSG] + state["messages"])
    return {
        "messages": [response],
        "input_file": state["input_file"]
    }

def controlled_gaia_tool_node(state: AgentState):
    # Find the last Human/User message
    for msg in reversed(state["messages"]):
        if isinstance(msg, HumanMessage):
            raw_user_message = msg.content
            break
        # Or if using HumanMessage directly:
        # if isinstance(msg, HumanMessage):
        #     raw_user_message = msg.content
        #     break
    else:
        raw_user_message = ""
    # Now, call the tool directly
    vector_store = get_gaia_data_vector_store(get_supabase_client())
    similar_question = vector_store.similarity_search(raw_user_message)
    # Format result for LLM
    example_msg = HumanMessage(
        content=f"Here I provide the most relevant question and answer from GAIA database: \n\n{similar_question[0].page_content}",
    )
    return {
        "messages": state["messages"] + [example_msg],   # Keep previous + append
        "input_file": state["input_file"]
    }


# # The graph
# builder = StateGraph(AgentState)

# # Define nodes: these do the work
# builder.add_node("assistant", assistant)
# builder.add_node("tools", ToolNode(tools))

# # Define edges: these determine how the control flow moves
# builder.add_edge(START, "assistant")
# builder.add_conditional_edges(
#     "assistant",
#     # If the latest message requires a tool, route to tools
#     # Otherwise, provide a direct response
#     tools_condition,
# )
# builder.add_edge("tools", "assistant")
# react_graph = builder.compile()

# Show the butler's thought process
# display(Image(react_graph.get_graph(xray=True).draw_mermaid_png()))

# Graph with retriever
builder = StateGraph(AgentState)
builder.add_node("retriever", controlled_gaia_tool_node)
builder.add_node("assistant", assistant)
builder.add_node("tools", ToolNode(tools))
builder.add_edge(START, "retriever")
builder.add_edge("retriever", "assistant")
builder.add_conditional_edges(
    "assistant",
    tools_condition,
)
builder.add_edge("tools", "assistant")

# Compile graph
react_graph = builder.compile()