from langgraph.graph import StateGraph
from src.langgraphagenticai.state.state import State
from langgraph.graph import START, END
from src.langgraphagenticai.nodes.basic_chatbot_node import BasicChatbotNode
from src.langgraphagenticai.tools.search_tool import get_tools, create_tool_node
from langgraph.prebuilt import tools_condition, ToolNode
from src.langgraphagenticai.nodes.chatbot_with_Tool_node import ChatbotWithToolNode
from src.langgraphagenticai.nodes.ai_news_node import AINewsNode


class GraphBuilder:
    def __init__(self, model):
        """
        Initializes the GraphBuilder class.

        Args:
            model: The language model (LLM) to be used in nodes.
        """
        # Store the provided language model
        self.llm = model

        # Initialize a StateGraph — the core structure used by LangGraph
        # to define workflows (nodes and their connections).
        self.graph_builder = StateGraph(State)

    # ----------------------------------------------------------------------
    def basic_chatbot_build_graph(self):
        """
        Builds a simple chatbot workflow graph.

        The graph consists of only one node (the chatbot node) that:
        - Takes a user message as input
        - Returns an AI-generated reply
        - Starts and ends at this single node
        """

        # Create an instance of the basic chatbot node using the LLM
        self.basic_chatbot_node = BasicChatbotNode(self.llm)

        # Add a node named "chatbot" and assign its processing function
        self.graph_builder.add_node("chatbot", self.basic_chatbot_node.process)

        # Define the flow of the graph:
        # Start → chatbot → End
        self.graph_builder.add_edge(START, "chatbot")
        self.graph_builder.add_edge("chatbot", END)

    # ----------------------------------------------------------------------
    def chatbot_with_tools_build_graph(self):
        """
        Builds an advanced chatbot workflow that can use external tools.

        This graph includes:
        - A chatbot node (main conversation handler)
        - A tools node (used when external info or web data is needed)
        - Conditional edges to decide when to call tools
        """

        # 1️⃣ Load the external tools (e.g., web search APIs)
        tools = get_tools()

        # 2️⃣ Create a LangGraph ToolNode using those tools
        tool_node = create_tool_node(tools)

        # 3️⃣ Initialize the LLM
        llm = self.llm

        # 4️⃣ Create a chatbot node that supports tool calling
        obj_chatbot_with_node = ChatbotWithToolNode(llm)
        chatbot_node = obj_chatbot_with_node.create_chatbot(tools)

        # 5️⃣ Add both chatbot and tool nodes to the graph
        self.graph_builder.add_node("chatbot", chatbot_node)
        self.graph_builder.add_node("tools", tool_node)

        # 6️⃣ Define how data flows between nodes:
        # Start → Chatbot
        self.graph_builder.add_edge(START, "chatbot")

        # From chatbot → tools (conditionally, only when the chatbot needs to call a tool)
        self.graph_builder.add_conditional_edges("chatbot", tools_condition)

        # From tools → chatbot (after tool execution, control returns to chatbot)
        self.graph_builder.add_edge("tools", "chatbot")

    # ----------------------------------------------------------------------
    def ai_news_builder_graph(self):
        """
        Builds a graph for the AI News use case.

        This workflow automates:
        1. Fetching AI-related news from the web (Tavily API)
        2. Summarizing it using an LLM
        3. Saving the result as a markdown file
        """

        # Create an AI news node object (handles fetch, summarize, and save)
        ai_news_node = AINewsNode(self.llm)

        # Add nodes for each stage of the process
        self.graph_builder.add_node("fetch_news", ai_news_node.fetch_news)
        self.graph_builder.add_node("summarize_news", ai_news_node.summarize_news)
        self.graph_builder.add_node("save_result", ai_news_node.save_result)

        # Define the flow of the process:
        # fetch_news → summarize_news → save_result → End
        self.graph_builder.set_entry_point("fetch_news")
        self.graph_builder.add_edge("fetch_news", "summarize_news")
        self.graph_builder.add_edge("summarize_news", "save_result")
        self.graph_builder.add_edge("save_result", END)

    # ----------------------------------------------------------------------
    def setup_graph(self, usecase: str):
        """
        Builds and compiles the appropriate LangGraph workflow
        based on the selected use case.

        Args:
            usecase (str): The name of the selected use case.
                           Options: "Basic Chatbot", "Chatbot With Web", "AI News"

        Returns:
            Compiled graph object ready for execution.
        """

        # Based on the selected use case, build the appropriate graph
        if usecase == "Basic Chatbot":
            self.basic_chatbot_build_graph()
        if usecase == "Chatbot With Web":
            self.chatbot_with_tools_build_graph()
        if usecase == "AI News":
            self.ai_news_builder_graph()

        # Compile the graph so it can be executed by LangGraph
        return self.graph_builder.compile()