from langgraph.graph import StateGraph, START, END
from .func import State, trim_history, call_practice_agent, call_teaching_agent
from langgraph.graph.state import CompiledStateGraph
from langgraph.checkpoint.memory import InMemorySaver


class LessonPractice2Agent:
    def __init__(self):
        pass

    @staticmethod
    def route_to_active_agent(state: State) -> str:
        if state["active_agent"] == "Practice Agent":
            return "Practice Agent"
        elif state["active_agent"] == "Teaching Agent":
            return "Teaching Agent"
        else:
            # Default to Teaching Agent if no active agent is set
            return "Teaching Agent"

    def node(self, graph: StateGraph):
        graph.add_node("trim_history", trim_history)
        graph.add_node("Practice Agent", call_practice_agent, destinations=("Teaching Agent",))
        graph.add_node(
            "Teaching Agent", call_teaching_agent, destinations=("Practice Agent",)
        )
        return graph

    def edge(self, graph: StateGraph):
        graph.add_edge(START, "trim_history")
        graph.add_conditional_edges(
            "trim_history",
            self.route_to_active_agent,
            {
                "Practice Agent": "Practice Agent",
                "Teaching Agent": "Teaching Agent",
            },
        )
        return graph

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


lesson_practice_agent = LessonPractice2Agent()