import copy
import json
from types import SimpleNamespace

import pytest
from lfx.graph import Graph
from lfx.graph.graph.utils import (
    find_last_node,
    process_flow,
    set_new_target_handle,
    ungroup_node,
    update_source_handle,
    update_target_handle,
    update_template,
)
from lfx.graph.vertex.base import Vertex
from lfx.interface.components import component_cache
from lfx.utils.flow_validation import CustomComponentValidationError

# Test cases for the graph module

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


@pytest.fixture
def sample_template():
    return {
        "field1": {"proxy": {"field": "some_field", "id": "node1"}},
        "field2": {"proxy": {"field": "other_field", "id": "node2"}},
    }


@pytest.fixture
def sample_nodes():
    return [
        {
            "id": "node1",
            "data": {"node": {"template": {"some_field": {"show": True, "advanced": False, "name": "Name1"}}}},
        },
        {
            "id": "node2",
            "data": {
                "node": {
                    "template": {
                        "other_field": {
                            "show": False,
                            "advanced": True,
                            "display_name": "DisplayName2",
                        }
                    }
                }
            },
        },
        {
            "id": "node3",
            "data": {"node": {"template": {"unrelated_field": {"show": True, "advanced": True}}}},
        },
    ]


def _settings_service(*, allow_custom_components: bool = False):
    return SimpleNamespace(
        settings=SimpleNamespace(
            allow_custom_components=allow_custom_components,
        )
    )


def _blocked_custom_flow() -> dict:
    return {
        "nodes": [
            {
                "id": "node-1",
                "data": {
                    "id": "node-1",
                    "type": "TotallyCustom",
                    "node": {
                        "display_name": "Blocked Node",
                        "template": {
                            "code": {"value": "print('blocked')"},
                        },
                    },
                },
            }
        ],
        "edges": [],
    }


def get_node_by_type(graph, node_type: type[Vertex]) -> Vertex | None:
    """Get a node by type."""
    return next((node for node in graph.vertices if isinstance(node, node_type)), None)


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


def test_from_payload_blocks_custom_components_when_disabled(monkeypatch):
    monkeypatch.setattr(
        "lfx.services.deps.get_settings_service",
        lambda: _settings_service(allow_custom_components=False),
    )
    monkeypatch.setattr(component_cache, "type_to_current_hash", {"ChatInput": {"known-hash"}})
    monkeypatch.setattr(component_cache, "all_types_dict", None)

    with pytest.raises(CustomComponentValidationError, match="custom components are not allowed"):
        Graph.from_payload(_blocked_custom_flow())


def test_find_last_node(grouped_chat_json_flow):
    grouped_chat_data = json.loads(grouped_chat_json_flow).get("data")
    nodes, edges = grouped_chat_data["nodes"], grouped_chat_data["edges"]
    last_node = find_last_node(nodes, edges)
    assert last_node is not None  # Replace with the actual expected value
    assert last_node["id"] == "LLMChain-pimAb"  # Replace with the actual expected value


def test_ungroup_node(grouped_chat_json_flow):
    grouped_chat_data = json.loads(grouped_chat_json_flow).get("data")
    group_node = grouped_chat_data["nodes"][2]  # Assuming the first node is a group node
    base_flow = copy.deepcopy(grouped_chat_data)
    ungroup_node(group_node["data"], base_flow)
    # after ungroup_node is called, the base_flow and grouped_chat_data should be different
    assert base_flow != grouped_chat_data
    # assert node 2 is not a group node anymore
    assert base_flow["nodes"][2]["data"]["node"].get("flow") is None
    # assert the edges are updated
    assert len(base_flow["edges"]) > len(grouped_chat_data["edges"])
    assert base_flow["edges"][0]["source"] == "ConversationBufferMemory-kUMif"
    assert base_flow["edges"][0]["target"] == "LLMChain-2P369"
    assert base_flow["edges"][1]["source"] == "PromptTemplate-Wjk4g"
    assert base_flow["edges"][1]["target"] == "LLMChain-2P369"
    assert base_flow["edges"][2]["source"] == "ChatOpenAI-rUJ1b"
    assert base_flow["edges"][2]["target"] == "LLMChain-2P369"


def test_process_flow(grouped_chat_json_flow):
    grouped_chat_data = json.loads(grouped_chat_json_flow).get("data")

    processed_flow = process_flow(grouped_chat_data)
    assert processed_flow is not None
    assert isinstance(processed_flow, dict)
    assert "nodes" in processed_flow
    assert "edges" in processed_flow


def test_process_flow_one_group(one_grouped_chat_json_flow):
    grouped_chat_data = json.loads(one_grouped_chat_json_flow).get("data")
    # There should be only one node
    assert len(grouped_chat_data["nodes"]) == 1
    # Get the node, it should be a group node
    group_node = grouped_chat_data["nodes"][0]
    node_data = group_node["data"]["node"]
    assert node_data.get("flow") is not None
    template_data = node_data["template"]
    assert any("openai_api_key" in key for key in template_data)
    # Get the openai_api_key dict
    openai_api_key = next(
        (template_data[key] for key in template_data if "openai_api_key" in key),
        None,
    )
    assert openai_api_key is not None
    assert openai_api_key["value"] == "test"

    processed_flow = process_flow(grouped_chat_data)
    assert processed_flow is not None
    assert isinstance(processed_flow, dict)
    assert "nodes" in processed_flow
    assert "edges" in processed_flow

    # Now get the node that has ChatOpenAI in its id
    chat_openai_node = next((node for node in processed_flow["nodes"] if "ChatOpenAI" in node["id"]), None)
    assert chat_openai_node is not None
    assert chat_openai_node["data"]["node"]["template"]["openai_api_key"]["value"] == "test"


def test_process_flow_vector_store_grouped(vector_store_grouped_json_flow):
    grouped_chat_data = json.loads(vector_store_grouped_json_flow).get("data")
    nodes = grouped_chat_data["nodes"]
    assert len(nodes) == 4
    # There are two group nodes in this flow
    # One of them is inside the other totalling 7 nodes
    # 4 nodes grouped, one of these turns into 1 normal node and 1 group node
    # This group node has 2 nodes inside it

    processed_flow = process_flow(grouped_chat_data)
    assert processed_flow is not None
    processed_nodes = processed_flow["nodes"]
    assert len(processed_nodes) == 7
    assert isinstance(processed_flow, dict)
    assert "nodes" in processed_flow
    assert "edges" in processed_flow
    edges = processed_flow["edges"]
    # Expected keywords in source and target fields
    expected_keywords = [
        {"source": "VectorStoreInfo", "target": "VectorStoreAgent"},
        {"source": "ChatOpenAI", "target": "VectorStoreAgent"},
        {"source": "OpenAIEmbeddings", "target": "Chroma"},
        {"source": "Chroma", "target": "VectorStoreInfo"},
        {"source": "WebBaseLoader", "target": "RecursiveCharacterTextSplitter"},
        {"source": "RecursiveCharacterTextSplitter", "target": "Chroma"},
    ]

    for idx, expected_keyword in enumerate(expected_keywords):
        for key, value in expected_keyword.items():
            assert value in edges[idx][key].split("-")[0], (
                f"Edge {idx}, key {key} expected to contain {value} but got {edges[idx][key]}"
            )


def test_update_template(sample_template, sample_nodes):
    # Making a deep copy to keep original sample_nodes unchanged
    nodes_copy = copy.deepcopy(sample_nodes)
    update_template(sample_template, nodes_copy)

    # Now, validate the updates.
    node1_updated = next((n for n in nodes_copy if n["id"] == "node1"), None)
    node2_updated = next((n for n in nodes_copy if n["id"] == "node2"), None)
    node3_updated = next((n for n in nodes_copy if n["id"] == "node3"), None)

    assert node1_updated["data"]["node"]["template"]["some_field"]["show"] is True
    assert node1_updated["data"]["node"]["template"]["some_field"]["advanced"] is False
    assert node1_updated["data"]["node"]["template"]["some_field"]["display_name"] == "Name1"

    assert node2_updated["data"]["node"]["template"]["other_field"]["show"] is False
    assert node2_updated["data"]["node"]["template"]["other_field"]["advanced"] is True
    assert node2_updated["data"]["node"]["template"]["other_field"]["display_name"] == "DisplayName2"

    # Ensure node3 remains unchanged
    assert node3_updated == sample_nodes[2]


# Test `update_target_handle`
def test_update_target_handle_proxy():
    new_edge = {
        "data": {
            "targetHandle": {
                "type": "some_type",
                "proxy": {"id": "some_id", "field": ""},
            }
        }
    }
    g_nodes = [{"id": "some_id", "data": {"node": {"flow": None}}}]
    updated_edge = update_target_handle(new_edge, g_nodes)
    assert updated_edge["data"]["targetHandle"] == new_edge["data"]["targetHandle"]


# Test `set_new_target_handle`
def test_set_new_target_handle():
    proxy_id = "proxy_id"
    new_edge = {"target": None, "data": {"targetHandle": {}}}
    target_handle = {"type": "type_1", "proxy": {"field": "field_1"}}
    node = {
        "data": {
            "node": {
                "flow": True,
                "template": {"field_1": {"proxy": {"field": "new_field", "id": "new_id"}}},
            }
        }
    }
    set_new_target_handle(proxy_id, new_edge, target_handle, node)
    assert new_edge["target"] == "proxy_id"
    assert new_edge["data"]["targetHandle"]["fieldName"] == "field_1"
    assert new_edge["data"]["targetHandle"]["proxy"] == {
        "field": "new_field",
        "id": "new_id",
    }


# Test `update_source_handle`
def test_update_source_handle():
    new_edge = {"source": None, "data": {"sourceHandle": {"id": None}}}
    flow_data = {
        "nodes": [{"id": "some_node"}, {"id": "last_node"}],
        "edges": [{"source": "some_node"}],
    }
    updated_edge = update_source_handle(new_edge, flow_data["nodes"], flow_data["edges"])
    assert updated_edge["source"] == "last_node"
    assert updated_edge["data"]["sourceHandle"]["id"] == "last_node"


def test_process_flow_grouped_loop_uses_output_proxy_when_internal_edges_are_cyclic():
    loop_id = "LoopComponent-loop"
    converter_id = "TypeConverterComponent-converter"
    group_id = "groupComponent-loop"
    sink_id = "ChatOutput-sink"
    flow_data = {
        "nodes": [
            {
                "id": group_id,
                "data": {
                    "id": group_id,
                    "type": "GroupNode",
                    "node": {
                        "template": {},
                        "outputs": [
                            {
                                "name": f"{loop_id}_done",
                                "proxy": {"id": loop_id, "name": "done", "nodeDisplayName": "Loop"},
                                "types": ["DataFrame"],
                            }
                        ],
                        "flow": {
                            "data": {
                                "nodes": [
                                    {
                                        "id": loop_id,
                                        "data": {
                                            "type": "LoopComponent",
                                            "node": {
                                                "template": {},
                                                "outputs": [
                                                    {"name": "item", "types": ["Data"], "allows_loop": True},
                                                    {"name": "done", "types": ["DataFrame"]},
                                                ],
                                            },
                                        },
                                    },
                                    {
                                        "id": converter_id,
                                        "data": {
                                            "type": "TypeConverterComponent",
                                            "node": {
                                                "template": {},
                                                "outputs": [{"name": "message_output", "types": ["Message"]}],
                                            },
                                        },
                                    },
                                ],
                                "edges": [
                                    {
                                        "source": loop_id,
                                        "target": converter_id,
                                        "data": {
                                            "sourceHandle": {
                                                "dataType": "LoopComponent",
                                                "id": loop_id,
                                                "name": "item",
                                                "output_types": ["Data"],
                                            },
                                            "targetHandle": {
                                                "fieldName": "input_data",
                                                "id": converter_id,
                                                "inputTypes": ["Data"],
                                                "type": "other",
                                            },
                                        },
                                    },
                                    {
                                        "source": converter_id,
                                        "target": loop_id,
                                        "data": {
                                            "sourceHandle": {
                                                "dataType": "TypeConverterComponent",
                                                "id": converter_id,
                                                "name": "message_output",
                                                "output_types": ["Message"],
                                            },
                                            "targetHandle": {
                                                "dataType": "LoopComponent",
                                                "id": loop_id,
                                                "name": "item",
                                                "output_types": ["Data", "Message"],
                                            },
                                        },
                                    },
                                ],
                            }
                        },
                    },
                },
            },
            {
                "id": sink_id,
                "data": {"type": "ChatOutput", "node": {"template": {}, "outputs": []}},
            },
        ],
        "edges": [
            {
                "source": group_id,
                "target": sink_id,
                "data": {
                    "sourceHandle": {
                        "dataType": "GroupNode",
                        "id": group_id,
                        "name": f"{loop_id}_done",
                        "output_types": ["DataFrame"],
                    },
                    "targetHandle": {
                        "fieldName": "input_value",
                        "id": sink_id,
                        "inputTypes": ["DataFrame"],
                        "type": "other",
                    },
                },
            }
        ],
    }

    processed_flow = process_flow(flow_data)

    updated_edge = next(edge for edge in processed_flow["edges"] if edge["target"] == sink_id)
    assert updated_edge["source"] == loop_id
    assert updated_edge["data"]["sourceHandle"] == {
        "dataType": "LoopComponent",
        "id": loop_id,
        "name": "done",
        "output_types": ["DataFrame"],
    }


def test_add_nodes_and_edges_recomputes_grouped_loop_cycles_after_processing(monkeypatch):
    loop_id = "LoopComponent-loop"
    converter_id = "TypeConverterComponent-converter"
    group_id = "groupComponent-loop"
    sink_id = "ChatOutput-sink"
    flow_data = {
        "nodes": [
            {
                "id": group_id,
                "data": {
                    "id": group_id,
                    "type": "GroupNode",
                    "node": {
                        "template": {},
                        "outputs": [
                            {
                                "name": f"{loop_id}_done",
                                "proxy": {"id": loop_id, "name": "done", "nodeDisplayName": "Loop"},
                            }
                        ],
                        "flow": {
                            "data": {
                                "nodes": [
                                    {"id": loop_id, "data": {"type": "LoopComponent", "node": {"template": {}}}},
                                    {
                                        "id": converter_id,
                                        "data": {"type": "TypeConverterComponent", "node": {"template": {}}},
                                    },
                                ],
                                "edges": [
                                    {
                                        "source": loop_id,
                                        "target": converter_id,
                                        "data": {
                                            "sourceHandle": {"id": loop_id, "name": "item"},
                                            "targetHandle": {"id": converter_id, "fieldName": "input_data"},
                                        },
                                    },
                                    {
                                        "source": converter_id,
                                        "target": loop_id,
                                        "data": {
                                            "sourceHandle": {"id": converter_id, "name": "message_output"},
                                            "targetHandle": {"id": loop_id, "name": "item"},
                                        },
                                    },
                                ],
                            }
                        },
                    },
                },
            },
            {"id": sink_id, "data": {"type": "ChatOutput", "node": {"template": {}}}},
        ],
        "edges": [
            {
                "source": group_id,
                "target": sink_id,
                "data": {
                    "sourceHandle": {"id": group_id, "name": f"{loop_id}_done"},
                    "targetHandle": {"id": sink_id, "fieldName": "input_value"},
                },
            }
        ],
    }

    monkeypatch.setattr(Graph, "initialize", lambda _: None)

    graph = Graph()
    graph.add_nodes_and_edges(flow_data["nodes"], flow_data["edges"])

    assert graph.cycle_vertices == {loop_id, converter_id}


# TODO: Move to Langflow tests
@pytest.mark.skip(reason="Temporarily disabled")
async def test_serialize_graph():
    pass
    # Get the actual starter projects and directly await the result
    # starter_projects = await load_starter_projects()
    # project_data = starter_projects[0][1]
    # data = project_data["data"]

    # # Create and test the graph
    # graph = Graph.from_payload(data)
    # assert isinstance(graph, Graph)
    # serialized = graph.dumps()
    # assert serialized is not None
    # assert isinstance(serialized, str)
    # assert len(serialized) > 0
