import json
from langflow.graph.nodes import (
    WrapperNode,
    AgentNode,
    ToolNode,
    ChainNode,
    PromptNode,
    LLMNode,
    ToolkitNode,
    FileToolNode,
)

import pytest
from langchain.agents import AgentExecutor
from langflow.graph import Edge, Graph, Node
from langflow.utils.payload import build_json, get_root_node


# Test cases for the graph module

# now we have three types of graph:
# BASIC_EXAMPLE_PATH, COMPLEX_EXAMPLE_PATH, OPENAPI_EXAMPLE_PATH


def get_graph(_type="basic"):
    """Get a graph from a json file"""
    if _type == "basic":
        path = pytest.BASIC_EXAMPLE_PATH
    elif _type == "complex":
        path = pytest.COMPLEX_EXAMPLE_PATH
    elif _type == "openapi":
        path = pytest.OPENAPI_EXAMPLE_PATH

    with open(path, "r") as f:
        flow_graph = json.load(f)
    data_graph = flow_graph["data"]
    nodes = data_graph["nodes"]
    edges = data_graph["edges"]
    return Graph(nodes, edges)


@pytest.fixture
def basic_graph():
    return get_graph()


@pytest.fixture
def complex_graph():
    return get_graph("complex")


@pytest.fixture
def openapi_graph():
    return get_graph("openapi")


def get_node_by_type(graph, node_type):
    """Get a node by type"""
    return next((node for node in graph.nodes if isinstance(node, node_type)), None)


def test_graph_structure(basic_graph):
    assert isinstance(basic_graph, Graph)
    assert len(basic_graph.nodes) > 0
    assert len(basic_graph.edges) > 0
    for node in basic_graph.nodes:
        assert isinstance(node, Node)
    for edge in basic_graph.edges:
        assert isinstance(edge, Edge)
        assert edge.source in basic_graph.nodes
        assert edge.target in basic_graph.nodes


def test_circular_dependencies(basic_graph):
    assert isinstance(basic_graph, Graph)

    def check_circular(node, visited):
        visited.add(node)
        neighbors = basic_graph.get_nodes_with_target(node)
        for neighbor in neighbors:
            if neighbor in visited:
                return True
            if check_circular(neighbor, visited.copy()):
                return True
        return False

    for node in basic_graph.nodes:
        assert not check_circular(node, set())


def test_invalid_node_types():
    graph_data = {
        "nodes": [
            {
                "id": "1",
                "data": {
                    "node": {
                        "base_classes": ["BaseClass"],
                        "template": {
                            "_type": "InvalidNodeType",
                        },
                    },
                },
            },
        ],
        "edges": [],
    }
    with pytest.raises(Exception):
        Graph(graph_data["nodes"], graph_data["edges"])


def test_get_nodes_with_target(basic_graph):
    """Test getting connected nodes"""
    assert isinstance(basic_graph, Graph)
    # Get root node
    root = get_root_node(basic_graph)
    assert root is not None
    connected_nodes = basic_graph.get_nodes_with_target(root)
    assert connected_nodes is not None


def test_get_node_neighbors_basic(basic_graph):
    """Test getting node neighbors"""

    assert isinstance(basic_graph, Graph)
    # Get root node
    root = get_root_node(basic_graph)
    assert root is not None
    neighbors = basic_graph.get_node_neighbors(root)
    assert neighbors is not None
    assert isinstance(neighbors, dict)
    # Root Node is an Agent, it requires an LLMChain and tools
    # We need to check if there is a Chain in the one of the neighbors'
    # data attribute in the type key
    assert any(
        "Chain" in neighbor.data["type"] for neighbor, val in neighbors.items() if val
    )
    # assert Serper Search is in the neighbors
    assert any(
        "Serper" in neighbor.data["type"] for neighbor, val in neighbors.items() if val
    )
    # Now on to the Chain's neighbors
    chain = next(
        neighbor
        for neighbor, val in neighbors.items()
        if "Chain" in neighbor.data["type"] and val
    )
    chain_neighbors = basic_graph.get_node_neighbors(chain)
    assert chain_neighbors is not None
    assert isinstance(chain_neighbors, dict)
    # Check if there is a LLM in the chain's neighbors
    assert any(
        "OpenAI" in neighbor.data["type"]
        for neighbor, val in chain_neighbors.items()
        if val
    )
    # Chain should have a Prompt as a neighbor
    assert any(
        "Prompt" in neighbor.data["type"]
        for neighbor, val in chain_neighbors.items()
        if val
    )


def test_get_node_neighbors_complex(complex_graph):
    """Test getting node neighbors"""
    assert isinstance(complex_graph, Graph)
    # Get root node
    root = get_root_node(complex_graph)
    assert root is not None
    neighbors = complex_graph.get_nodes_with_target(root)
    assert neighbors is not None
    # Neighbors should be a list of nodes
    assert isinstance(neighbors, list)
    # Root Node is an Agent, it requires an LLMChain and tools
    # We need to check if there is a Chain in the one of the neighbors'
    assert any("Chain" in neighbor.data["type"] for neighbor in neighbors)
    # assert Tool is in the neighbors
    assert any("Tool" in neighbor.data["type"] for neighbor in neighbors)
    # Now on to the Chain's neighbors
    chain = next(neighbor for neighbor in neighbors if "Chain" in neighbor.data["type"])
    chain_neighbors = complex_graph.get_nodes_with_target(chain)
    assert chain_neighbors is not None
    # Check if there is a LLM in the chain's neighbors
    assert any("OpenAI" in neighbor.data["type"] for neighbor in chain_neighbors)
    # Chain should have a Prompt as a neighbor
    assert any("Prompt" in neighbor.data["type"] for neighbor in chain_neighbors)
    # Now on to the Tool's neighbors
    tool = next(neighbor for neighbor in neighbors if "Tool" in neighbor.data["type"])
    tool_neighbors = complex_graph.get_nodes_with_target(tool)
    assert tool_neighbors is not None
    # Check if there is an Agent in the tool's neighbors
    assert any("Agent" in neighbor.data["type"] for neighbor in tool_neighbors)
    # This Agent has a Tool that has a PythonFunction as func
    agent = next(
        neighbor for neighbor in tool_neighbors if "Agent" in neighbor.data["type"]
    )
    agent_neighbors = complex_graph.get_nodes_with_target(agent)
    assert agent_neighbors is not None
    # Check if there is a Tool in the agent's neighbors
    assert any("Tool" in neighbor.data["type"] for neighbor in agent_neighbors)
    # This Tool has a PythonFunction as func
    tool = next(
        neighbor for neighbor in agent_neighbors if "Tool" in neighbor.data["type"]
    )
    tool_neighbors = complex_graph.get_nodes_with_target(tool)
    assert tool_neighbors is not None
    # Check if there is a PythonFunction in the tool's neighbors
    assert any("PythonFunction" in neighbor.data["type"] for neighbor in tool_neighbors)


def test_get_node(basic_graph):
    """Test getting a single node"""
    node_id = basic_graph.nodes[0].id
    node = basic_graph.get_node(node_id)
    assert isinstance(node, Node)
    assert node.id == node_id


def test_build_nodes(basic_graph):
    """Test building nodes"""

    assert len(basic_graph.nodes) == len(basic_graph._nodes)
    for node in basic_graph.nodes:
        assert isinstance(node, Node)


def test_build_edges(basic_graph):
    """Test building edges"""
    assert len(basic_graph.edges) == len(basic_graph._edges)
    for edge in basic_graph.edges:
        assert isinstance(edge, Edge)
        assert isinstance(edge.source, Node)
        assert isinstance(edge.target, Node)


def test_get_root_node(basic_graph, complex_graph):
    """Test getting root node"""
    assert isinstance(basic_graph, Graph)
    root = get_root_node(basic_graph)
    assert root is not None
    assert isinstance(root, Node)
    assert root.data["type"] == "ZeroShotAgent"
    # For complex example, the root node is a ZeroShotAgent too
    assert isinstance(complex_graph, Graph)
    root = get_root_node(complex_graph)
    assert root is not None
    assert isinstance(root, Node)
    assert root.data["type"] == "ZeroShotAgent"


def test_build_json(basic_graph):
    """Test building JSON from graph"""
    assert isinstance(basic_graph, Graph)
    root = get_root_node(basic_graph)
    json_data = build_json(root, basic_graph)
    assert isinstance(json_data, dict)
    assert json_data["_type"] == "zero-shot-react-description"
    assert isinstance(json_data["llm_chain"], dict)
    assert json_data["llm_chain"]["_type"] == "llm_chain"
    assert json_data["llm_chain"]["memory"] is None
    assert json_data["llm_chain"]["verbose"] is False
    assert isinstance(json_data["llm_chain"]["prompt"], dict)
    assert isinstance(json_data["llm_chain"]["llm"], dict)
    assert json_data["llm_chain"]["output_key"] == "text"
    assert isinstance(json_data["allowed_tools"], list)
    assert all(isinstance(tool, dict) for tool in json_data["allowed_tools"])
    assert isinstance(json_data["return_values"], list)
    assert all(isinstance(val, str) for val in json_data["return_values"])


def test_validate_edges(basic_graph):
    """Test validating edges"""

    assert isinstance(basic_graph, Graph)
    # all edges should be valid
    assert all(edge.valid for edge in basic_graph.edges)


def test_matched_type(basic_graph):
    """Test matched type attribute in Edge"""
    assert isinstance(basic_graph, Graph)
    # all edges should be valid
    assert all(edge.valid for edge in basic_graph.edges)
    # all edges should have a matched_type attribute
    assert all(hasattr(edge, "matched_type") for edge in basic_graph.edges)
    # The matched_type attribute should be in the source_types attr
    assert all(edge.matched_type in edge.source_types for edge in basic_graph.edges)


def test_build_params(basic_graph):
    """Test building params"""

    assert isinstance(basic_graph, Graph)
    # all edges should be valid
    assert all(edge.valid for edge in basic_graph.edges)
    # all edges should have a matched_type attribute
    assert all(hasattr(edge, "matched_type") for edge in basic_graph.edges)
    # The matched_type attribute should be in the source_types attr
    assert all(edge.matched_type in edge.source_types for edge in basic_graph.edges)
    # Get the root node
    root = get_root_node(basic_graph)
    # Root node is a ZeroShotAgent
    # which requires an llm_chain, allowed_tools and return_values
    assert isinstance(root.params, dict)
    assert "llm_chain" in root.params
    assert "allowed_tools" in root.params
    assert "return_values" in root.params
    # The llm_chain should be a Node
    assert isinstance(root.params["llm_chain"], Node)
    # The allowed_tools should be a list of Nodes
    assert isinstance(root.params["allowed_tools"], list)
    assert all(isinstance(tool, Node) for tool in root.params["allowed_tools"])
    # The return_values is of type str so it should be a list of strings
    assert isinstance(root.params["return_values"], list)
    assert all(isinstance(val, str) for val in root.params["return_values"])
    # The llm_chain should have a prompt and llm
    llm_chain_node = root.params["llm_chain"]
    assert isinstance(llm_chain_node.params, dict)
    assert "prompt" in llm_chain_node.params
    assert "llm" in llm_chain_node.params
    # The prompt should be a Node
    assert isinstance(llm_chain_node.params["prompt"], Node)
    # The llm should be a Node
    assert isinstance(llm_chain_node.params["llm"], Node)
    # The prompt should have format_insctructions, suffix, prefix
    prompt_node = llm_chain_node.params["prompt"]
    assert isinstance(prompt_node.params, dict)
    assert "format_instructions" in prompt_node.params
    assert "suffix" in prompt_node.params
    assert "prefix" in prompt_node.params
    # All of them should be of type str
    assert isinstance(prompt_node.params["format_instructions"], str)
    assert isinstance(prompt_node.params["suffix"], str)
    assert isinstance(prompt_node.params["prefix"], str)
    # The llm should have a model
    llm_node = llm_chain_node.params["llm"]
    assert isinstance(llm_node.params, dict)
    assert "model_name" in llm_node.params
    # The model should be a str
    assert isinstance(llm_node.params["model_name"], str)


def test_build(basic_graph, complex_graph):
    """Test Node's build method"""
    # def build(self):
    # # The params dict is used to build the module
    # # it contains values and keys that point to nodes which
    # # have their own params dict
    # # When build is called, we iterate through the params dict
    # # and if the value is a node, we call build on that node
    # # and use the output of that build as the value for the param
    # # if the value is not a node, then we use the value as the param
    # # and continue
    # # Another aspect is that the node_type is the class that we need to import
    # # and instantiate with these built params

    # # Build each node in the params dict
    # for key, value in self.params.items():
    #     if isinstance(value, Node):
    #         self.params[key] = value.build()

    # # Get the class from LANGCHAIN_TYPES_DICT
    # # and instantiate it with the params
    # # and return the instance
    # return LANGCHAIN_TYPES_DICT[self.node_type](**self.params)

    assert isinstance(basic_graph, Graph)
    # Now we test the build method
    # Build the Agent
    agent = basic_graph.build()
    # The agent should be a AgentExecutor
    assert isinstance(agent, AgentExecutor)

    # Now we test the complex example
    assert isinstance(complex_graph, Graph)
    # Now we test the build method
    agent = complex_graph.build()
    # The agent should be a AgentExecutor
    assert isinstance(agent, AgentExecutor)


def test_agent_node_build(basic_graph):
    agent_node = get_node_by_type(basic_graph, AgentNode)
    assert agent_node is not None
    built_object = agent_node.build()
    assert built_object is not None
    # Add any further assertions specific to the AgentNode's build() method


def test_tool_node_build(basic_graph):
    tool_node = get_node_by_type(basic_graph, ToolNode)
    assert tool_node is not None
    built_object = tool_node.build()
    assert built_object is not None
    # Add any further assertions specific to the ToolNode's build() method


def test_chain_node_build(complex_graph):
    chain_node = get_node_by_type(complex_graph, ChainNode)
    assert chain_node is not None
    built_object = chain_node.build()
    assert built_object is not None
    # Add any further assertions specific to the ChainNode's build() method


def test_prompt_node_build(complex_graph):
    prompt_node = get_node_by_type(complex_graph, PromptNode)
    assert prompt_node is not None
    built_object = prompt_node.build()
    assert built_object is not None
    # Add any further assertions specific to the PromptNode's build() method


def test_llm_node_build(basic_graph):
    llm_node = get_node_by_type(basic_graph, LLMNode)
    assert llm_node is not None
    built_object = llm_node.build()
    assert built_object is not None
    # Add any further assertions specific to the LLMNode's build() method


def test_toolkit_node_build(openapi_graph):
    toolkit_node = get_node_by_type(openapi_graph, ToolkitNode)
    assert toolkit_node is not None
    built_object = toolkit_node.build()
    assert built_object is not None
    # Add any further assertions specific to the ToolkitNode's build() method


def test_file_tool_node_build(openapi_graph):
    file_tool_node = get_node_by_type(openapi_graph, FileToolNode)
    assert file_tool_node is not None
    built_object = file_tool_node.build()
    assert built_object is not None
    # Add any further assertions specific to the FileToolNode's build() method


def test_wrapper_node_build(openapi_graph):
    wrapper_node = get_node_by_type(openapi_graph, WrapperNode)
    assert wrapper_node is not None
    built_object = wrapper_node.build()
    assert built_object is not None
    # Add any further assertions specific to the WrapperNode's build() method