import os
from typing import Annotated, TypedDict, Sequence
from langgraph.graph import StateGraph, START, END
from langchain_core.messages import BaseMessage, SystemMessage, HumanMessage, AIMessage
from langgraph.graph.message import add_messages
from langchain_google_genai import ChatGoogleGenerativeAI as Gemini
from langgraph.prebuilt import ToolNode
from tools import tools
from vector_db import vector_db, add_recent_question
from dotenv import load_dotenv


load_dotenv()

# initiate model
api_key = os.getenv("gemini_api_key")
model = Gemini(model="gemini-2.0-flash", temperature=0.8, api_key=api_key).bind_tools(
    tools=tools
)


class MessageState(TypedDict):
    messages: Annotated[Sequence[BaseMessage], add_messages]


# creating nodes


def retrieve(state: MessageState) -> MessageState:
    user_input = state["messages"][-1].content
    similar_questions = vector_db.similarity_search(user_input, k=1)
    if similar_questions:
        reference_msg = HumanMessage(
            content=f"Here is a similar question and answer from history:\n\n{similar_questions[0].page_content}"
        )
    else:
        reference_msg = HumanMessage(
            content="No similar question found in the database."
        )

    add_recent_question(user_input)
    return {"messages": [reference_msg]}


def call_model(state: MessageState) -> MessageState:
    response = model.invoke(state["messages"])
    return {"messages": [response]}


def should_continue(state: MessageState) -> str:
    last_message = state["messages"][-1]
    if isinstance(last_message, AIMessage) and last_message.tool_calls:
        return "tools_node"
    return "end"


tools_node = ToolNode(tools=tools)

graph = StateGraph(MessageState)

graph.add_node("retriever_node", retrieve)
graph.add_node("model_node", call_model)
graph.add_node("tools_node", tools_node)

graph.add_edge(START, "retriever_node")
graph.add_edge("retriever_node", "model_node")
graph.add_conditional_edges(
    "model_node",
    should_continue,
    {
        "tools_node": "tools_node",
        "end": END,
    },
)
graph.add_edge("tools_node", "model_node")
agent = graph.compile()
#testing
if __name__ == "__main__":
    with open("./system_prompt.txt", "r", encoding="utf-8") as f:
        SYSTEM_PROMPT_CONTENT = f.read()
        SYSTEM_MESSAGE = SystemMessage(content=SYSTEM_PROMPT_CONTENT)
        user_input = ""
        while user_input != "exit":
            user_input = input("question : ")
            user_query = HumanMessage(content=user_input)
            response = agent.invoke({"messages": [SYSTEM_MESSAGE] + [user_query]})
            last_message = response["messages"][-1]
            if isinstance(last_message, AIMessage):
                print(f"Agent: {last_message.content[14:]}")
                print(f"AI-debug: {last_message}")
            else:
                print(f"AI-debug: {last_message}")