from langgraph.graph import StateGraph, START, END
from .func import State
from langgraph.graph.state import CompiledStateGraph
from langgraph.store.memory import InMemoryStore


class PrimaryChatBot:
    def __init__(self):
        pass

    @staticmethod
    def should_continue(state: State):
        messages = state["messages"]
        last_message = messages[-1]
        if not last_message.tool_calls:
            return "end"
        else:
            return "continue"

    def node(self, graph: StateGraph):

        return graph

    def edge(self, graph: StateGraph):

        return graph

    def __call__(self, checkpointer=InMemoryStore()) -> CompiledStateGraph:
        graph = StateGraph(State)
        graph: StateGraph = self.node(graph)
        graph: StateGraph = self.edge(graph)
        return graph.compile(checkpointer=checkpointer)