from langgraph.graph import StateGraph, START, END
from .func import (
    State,
    trim_history,
    extract_transcript_and_comment,
    script_structure_analyzer,
    comment_insight_extractor,
    scientific_fact_finder,
    script_re_outline,
    script_writer_init,
    script_writer_single,
    should_continue_writing,
    script_writer_end,
)
from langgraph.graph.state import CompiledStateGraph


class AgentTranscript:
    def __init__(self):
        self.builder = StateGraph(State)

    @staticmethod
    def check_existed_comment(state: State):
        if not state["comment"]:
            return "scientific_fact_finder"
        else:
            return "comment_insight_extractor"

    def node(self):
        self.builder.add_node("trim_history", trim_history)
        self.builder.add_node(
            "extract_transcript_and_comment", extract_transcript_and_comment
        )
        self.builder.add_node("script_structure_analyzer", script_structure_analyzer)
        self.builder.add_node("comment_insight_extractor", comment_insight_extractor)
        self.builder.add_node("scientific_fact_finder", scientific_fact_finder)
        self.builder.add_node("script_re_outline", script_re_outline)
        self.builder.add_node("script_writer_init", script_writer_init)
        self.builder.add_node("script_writer_single", script_writer_single)
        self.builder.add_node("script_writer_end", script_writer_end)

    def edge(self):
        self.builder.add_edge(START, "trim_history")
        self.builder.add_edge("trim_history", "extract_transcript_and_comment")
        self.builder.add_edge(
            "extract_transcript_and_comment", "script_structure_analyzer"
        )
        self.builder.add_conditional_edges(
            "script_structure_analyzer",
            self.check_existed_comment,
            {
                "scientific_fact_finder": "scientific_fact_finder",
                "comment_insight_extractor": "comment_insight_extractor",
            },
        )
        self.builder.add_edge("comment_insight_extractor", "scientific_fact_finder")
        self.builder.add_edge("scientific_fact_finder", "script_re_outline")
        self.builder.add_edge("script_re_outline", "script_writer_init")

        # Conditional routing for script writing
        self.builder.add_conditional_edges(
            "script_writer_init",
            should_continue_writing,
            {
                "script_writer_single": "script_writer_single",
                "script_writer_end": "script_writer_end",
            },
        )

        self.builder.add_conditional_edges(
            "script_writer_single",
            should_continue_writing,
            {
                "script_writer_single": "script_writer_single",
                "script_writer_end": "script_writer_end",
            },
        )

        self.builder.add_edge("script_writer_end", END)

    def __call__(self) -> CompiledStateGraph:
        self.node()
        self.edge()

        return self.builder.compile()


script_writer_agent = AgentTranscript()()