# graph_manager.py
# graph definition and compilation

from langgraph.graph import StateGraph, START, END
from langgraph.checkpoint.memory import MemorySaver
from nodes import (
    State,
    speaker_chatbot_node,
    tool_node
)
from langgraph.prebuilt import tools_condition
from form_management.form_management import Form
import sys
from IPython.display import Image, display
import logging

class GraphManager:
    def __init__(self):
        self.graph = None
        self.init_graph()

    # user input -> 
    #   if user intent is clear -> speaker form management -> audit review and feedback -> speaker generate response based on feedback
    #   if user intent is unclear -> speaker generate response for clarification
    def init_graph(self):
        graph_memory = MemorySaver()
        graph_builder = StateGraph(State)

        graph_builder.add_node("speaker_chatbot", speaker_chatbot_node)
        graph_builder.add_node("tools", tool_node)

        graph_builder.add_conditional_edges(
            "speaker_chatbot",
            tools_condition,
        )
        # Any time a tool is called, we return to the chatbot to decide the next step
        graph_builder.add_edge("tools", "speaker_chatbot")

        graph_builder.add_edge(START, "speaker_chatbot")
        graph_builder.add_edge("speaker_chatbot", END)

        self.graph = graph_builder.compile(checkpointer=graph_memory)

    def save_graph_image(self):
        """
        Save the graph as a PNG image.
        """
        image = Image(self.graph.get_graph().draw_mermaid_png())
        with open("graph.png", "wb") as f:
            f.write(image.data)

    def process_user_input(self, user_input, session_id):
        """
        Process the user input and update the graph state.
        """
        config = {"configurable": {"thread_id": session_id}}

        input_dict = {
            "messages": [("user", user_input)],
            "session_id": session_id
        }

        for event in self.graph.stream(
            input_dict,
            config,
            stream_mode="values"
        ):
            event["messages"][-1].pretty_print()

        snapshot = self.graph.get_state(config)
        return snapshot[0]["messages"][-1].content
    
def user_input_handler(user_input, session_id):
    logging.info(user_input)
    response = graph_manager.process_user_input(user_input, session_id)
    return response

graph_manager = GraphManager()
graph_manager.save_graph_image()

# # Example usage
# response = graph_manager.process_user_input(
#     "hi",
#     1
# )

# response = graph_manager.process_user_input(
#     r"C:\Users\Gavin\Desktop\hello-earth\data\jpeg_raw_data\bill\กองทุนพัฒนาบ้านห้วยทราย อ.แม่ออน จ.เชียงใหม่_page-0005.jpg",
#     1
# )

# response = graph_manager.process_user_input(
#     r"C:\Users\Gavin\Desktop\hello-earth\data\jpeg_raw_data\bill\กองทุนพัฒนาบ้านห้วยทราย อ.แม่ออน จ.เชียงใหม่_page-0004.jpg",
#     1
# )

# response = graph_manager.process_user_input(
#     "For the irrigation systems.",
#     1
# )

# response = graph_manager.process_user_input(
#     "My name is actually Oya Punpun. Please correct this.",
#     1
# )

# response = graph_manager.process_user_input(
#     "My middle name is Baba.",
#     1
# )

# response = graph_manager.process_user_input(
#     r"C:\Users\Gavin\Desktop\hello-earth\data\jpeg_raw_data\bill\กองทุนพัฒนาบ้านห้วยทราย อ.แม่ออน จ.เชียงใหม่_page-0004.jpg",
#     1
# )