from uuid import UUID
from typing import Any, Dict, List, cast

from vellum.workflows.constants import VellumIntegrationProviderType
from vellum.workflows.graph import Graph
from vellum.workflows.inputs import BaseInputs
from vellum.workflows.nodes import BaseNode, InlineSubworkflowNode
from vellum.workflows.outputs.base import BaseOutputs
from vellum.workflows.ports.port import Port
from vellum.workflows.references.lazy import LazyReference
from vellum.workflows.state import BaseState
from vellum.workflows.triggers.integration import IntegrationTrigger
from vellum.workflows.workflows.base import BaseWorkflow
from vellum_ee.workflows.display.base import EdgeDisplay, WorkflowInputsDisplay
from vellum_ee.workflows.display.editor.types import NodeDisplayData, NodeDisplayPosition
from vellum_ee.workflows.display.nodes import BaseNodeDisplay
from vellum_ee.workflows.display.workflows.base_workflow_display import BaseWorkflowDisplay
from vellum_ee.workflows.display.workflows.get_vellum_workflow_display_class import get_workflow_display


def test_base_workflow_display__serialize_empty_workflow():
    # GIVEN an empty workflow
    class ExampleWorkflow(BaseWorkflow):
        pass

    display = get_workflow_display(workflow_class=ExampleWorkflow)

    # WHEN serializing the workflow
    exec_config = display.serialize()

    # THEN it should return the expected config
    assert exec_config == {
        "input_variables": [],
        "state_variables": [],
        "output_variables": [],
        "workflow_raw_data": {
            "definition": {
                "module": ["vellum_ee", "workflows", "display", "tests", "test_base_workflow_display"],
                "name": "ExampleWorkflow",
            },
            "display_data": {"viewport": {"x": 0.0, "y": 0.0, "zoom": 1.0}},
            "edges": [],
            "nodes": [
                {
                    "data": {"label": "Entrypoint Node", "source_handle_id": "0af025a4-3b25-457d-a7ae-e3a7ba15c86c"},
                    "base": None,
                    "definition": None,
                    "display_data": {"position": {"x": 0.0, "y": -50.0}},
                    "id": "3c41cdd9-999a-48b8-9088-f6dfa1369bfd",
                    "inputs": [],
                    "type": "ENTRYPOINT",
                }
            ],
            "output_values": [],
        },
    }


def test_vellum_workflow_display__serialize_input_variables_with_capitalized_variable_override():
    # GIVEN a workflow with input variables
    class Inputs(BaseInputs):
        foo: str

    class StartNode(BaseNode):
        class Outputs(BaseNode.Outputs):
            output = Inputs.foo

    class ExampleWorkflow(BaseWorkflow[Inputs, BaseState]):
        graph = StartNode

    class ExampleWorkflowDisplay(BaseWorkflowDisplay[ExampleWorkflow]):
        inputs_display = {
            Inputs.foo: WorkflowInputsDisplay(id=UUID("97b63d71-5413-417f-9cf5-49e1b4fd56e4"), name="Foo")
        }

    display = get_workflow_display(
        base_display_class=ExampleWorkflowDisplay,
        workflow_class=ExampleWorkflow,
    )

    # WHEN serializing the workflow
    exec_config = display.serialize()

    # THEN the input variables are what we expect
    input_variables = exec_config["input_variables"]

    assert input_variables == [
        {
            "id": "97b63d71-5413-417f-9cf5-49e1b4fd56e4",
            "key": "Foo",
            "type": "STRING",
            "default": None,
            "required": True,
            "extensions": {"color": None},
        }
    ]


def test_vellum_workflow_display_serialize_valid_handle_ids_for_base_nodes():
    # GIVEN a workflow between two base nodes
    class StartNode(BaseNode):
        pass

    class EndNode(BaseNode):
        class Outputs(BaseNode.Outputs):
            hello = "world"

    class Workflow(BaseWorkflow):
        graph = StartNode >> EndNode

        class Outputs(BaseWorkflow.Outputs):
            final_value = EndNode.Outputs.hello

    # AND a display class for this workflow
    workflow_display = get_workflow_display(workflow_class=Workflow)

    # WHEN we serialize the workflow
    exec_config = workflow_display.serialize()

    # THEN the serialized workflow handle ids are valid
    raw_data = exec_config.get("workflow_raw_data")
    assert isinstance(raw_data, dict)
    nodes = raw_data.get("nodes")
    edges = raw_data.get("edges")

    assert isinstance(nodes, list)
    assert isinstance(edges, list)

    edge_source_handle_ids = {edge.get("source_handle_id") for edge in edges if isinstance(edge, dict)}
    edge_target_handle_ids = {edge.get("target_handle_id") for edge in edges if isinstance(edge, dict)}

    for node in nodes:
        assert isinstance(node, dict)

        if node["type"] in {"ENTRYPOINT", "TERMINAL"}:
            continue

        ports = node.get("ports")
        assert isinstance(ports, list)
        for port in ports:
            assert isinstance(port, dict)
            assert (
                port["id"] in edge_source_handle_ids
            ), f"Port {port['id']} from node {node['label']} not found in edge source handle ids"

        assert isinstance(node["trigger"], dict)
        assert (
            node["trigger"]["id"] in edge_target_handle_ids
        ), f"Trigger {node['trigger']['id']} from node {node['label']} not found in edge target handle ids"


def test_vellum_workflow_display__serialize_with_unused_nodes_and_edges():
    # GIVEN a workflow with active and unused nodes
    class NodeA(BaseNode):
        class Outputs(BaseNode.Outputs):
            result: str

    class NodeB(BaseNode):
        pass

    class NodeC(BaseNode):
        pass

    # AND A workflow that uses them correctly
    class Workflow(BaseWorkflow):
        graph = NodeA
        unused_graphs = {NodeB >> NodeC}

        class Outputs(BaseWorkflow.Outputs):
            final = NodeA.Outputs.result

    # WHEN we serialize it
    workflow_display = get_workflow_display(workflow_class=Workflow)

    # WHEN we serialize the workflow
    exec_config = workflow_display.serialize()

    # THEN the serialized workflow contains the expected nodes and edges
    raw_data = exec_config["workflow_raw_data"]
    assert isinstance(raw_data, dict)

    nodes = raw_data["nodes"]
    edges = raw_data["edges"]

    assert isinstance(nodes, list)
    assert isinstance(edges, list)

    # Find nodes by their definition name
    node_ids: Dict[str, str] = {}

    for node in nodes:
        assert isinstance(node, dict)
        definition = node.get("definition")
        if definition is None:
            continue

        assert isinstance(definition, dict)
        name = definition.get("name")
        if not isinstance(name, str):
            continue

        if name in ["NodeA", "NodeB", "NodeC"]:
            node_id = node.get("id")
            if isinstance(node_id, str):
                node_ids[name] = node_id

    # Verify all nodes are present
    assert "NodeA" in node_ids, "Active node NodeA not found in serialized output"
    assert "NodeB" in node_ids, "Unused node NodeB not found in serialized output"
    assert "NodeC" in node_ids, "Unused node NodeC not found in serialized output"

    # Verify the edge between NodeB and NodeC is present
    edge_found = False
    for edge in edges:
        assert isinstance(edge, dict)
        source_id = edge.get("source_node_id")
        target_id = edge.get("target_node_id")

        if (
            isinstance(source_id, str)
            and isinstance(target_id, str)
            and source_id == node_ids["NodeB"]
            and target_id == node_ids["NodeC"]
        ):
            edge_found = True
            break

    assert edge_found, "Edge between unused nodes NodeB and NodeC not found in serialized output"


def test_vellum_workflow_display__serialize_with_parse_json_expression():
    # GIVEN a workflow that uses the parse_json function
    from vellum.workflows.references.constant import ConstantValueReference

    class JsonNode(BaseNode):

        class Outputs(BaseNode.Outputs):
            json_result = ConstantValueReference('{"key": "value"}').parse_json()

    class Workflow(BaseWorkflow):
        graph = JsonNode

        class Outputs(BaseWorkflow.Outputs):
            final = JsonNode.Outputs.json_result

    # AND a display class for this workflow
    workflow_display = get_workflow_display(workflow_class=Workflow)

    # WHEN we serialize the workflow
    exec_config = workflow_display.serialize()

    # THEN the serialized workflow contains the parse_json expression
    raw_data = exec_config["workflow_raw_data"]
    assert isinstance(raw_data, dict)

    nodes = raw_data["nodes"]
    assert isinstance(nodes, list)

    json_node = None
    for node in nodes:
        assert isinstance(node, dict)
        definition = node.get("definition")
        if node.get("type") == "GENERIC" and isinstance(definition, dict) and definition.get("name") == "JsonNode":
            json_node = node
            break

    assert json_node is not None

    outputs = json_node.get("outputs", [])
    assert isinstance(outputs, list)

    json_result = None
    for output in outputs:
        assert isinstance(output, dict)
        if output.get("name") == "json_result":
            json_result = output
            break

    assert json_result == {
        "id": "e73fd6b1-1109-4a97-8510-c9ba8e6f5dbe",
        "name": "json_result",
        "type": "JSON",
        "value": {
            "type": "UNARY_EXPRESSION",
            "lhs": {
                "type": "CONSTANT_VALUE",
                "value": {
                    "type": "STRING",
                    "value": '{"key": "value"}',
                },
            },
            "operator": "parseJson",
        },
    }


def test_serialize__port_with_lazy_reference():
    # GIVEN a node with a lazy reference in a Port
    class MyNode(BaseNode):

        class Ports(BaseNode.Ports):
            foo = Port.on_if(LazyReference(lambda: MyNode.Outputs.bar))

        class Outputs(BaseNode.Outputs):
            bar: bool

    # AND a workflow that uses the node
    class Workflow(BaseWorkflow):
        graph = MyNode

    # WHEN we serialize the workflow
    workflow_display = get_workflow_display(workflow_class=Workflow)
    exec_config = workflow_display.serialize()

    # THEN the lazy reference should be serialized correctly
    raw_data = exec_config["workflow_raw_data"]
    assert isinstance(raw_data, dict)

    nodes = raw_data["nodes"]
    assert isinstance(nodes, list)

    my_node = nodes[1]
    assert isinstance(my_node, dict)
    ports = my_node.get("ports")
    assert isinstance(ports, list)
    assert ports == [
        {
            "id": "02e90580-e6b6-441c-a5d4-7cd083c11fc7",
            "name": "foo",
            "type": "IF",
            "expression": {
                "type": "NODE_OUTPUT",
                "node_id": str(MyNode.__id__),
                "node_output_id": str(MyNode.__output_ids__["bar"]),
            },
        }
    ]


def test_global_propagation_deep_nested_subworkflows():
    # GIVEN the root workflow, a middle workflow, and an inner workflow

    class RootInputs(BaseInputs):
        root_param: str

    class MiddleInputs(BaseInputs):
        middle_param: str

    class InnerInputs(BaseInputs):
        inner_param: str

    class InnerNode(BaseNode):
        class Outputs(BaseOutputs):
            done: bool

        def run(self) -> Outputs:
            return self.Outputs(done=True)

    class InnerWorkflow(BaseWorkflow[InnerInputs, BaseState]):
        graph = InnerNode

    class MiddleInlineSubworkflowNode(InlineSubworkflowNode):
        subworkflow_inputs = {"inner_param": "x"}
        subworkflow = InnerWorkflow

    class MiddleWorkflow(BaseWorkflow[MiddleInputs, BaseState]):
        graph = MiddleInlineSubworkflowNode

    class OuterInlineSubworkflowNode(InlineSubworkflowNode):
        subworkflow_inputs = {"middle_param": "y"}
        subworkflow = MiddleWorkflow

    class RootWorkflow(BaseWorkflow[RootInputs, BaseState]):
        graph = OuterInlineSubworkflowNode

    # WHEN we build the displays
    root_display = get_workflow_display(workflow_class=RootWorkflow)
    middle_display = get_workflow_display(
        workflow_class=MiddleWorkflow, parent_display_context=root_display.display_context
    )
    inner_display = get_workflow_display(
        workflow_class=InnerWorkflow, parent_display_context=middle_display.display_context
    )

    # THEN the deepest display must include root + middle + inner inputs in its GLOBAL view
    inner_global_names = {ref.name for ref in inner_display.display_context.global_workflow_input_displays.keys()}

    assert inner_global_names == {"middle_param", "inner_param", "root_param"}


def test_serialize_workflow_with_edge_display_data():
    """
    Tests that edges with z_index values serialize display_data correctly.
    """

    # GIVEN a workflow with connected nodes
    class StartNode(BaseNode):
        class Outputs(BaseNode.Outputs):
            result: str

    class EndNode(BaseNode):
        class Outputs(BaseNode.Outputs):
            final: str

    class TestWorkflow(BaseWorkflow):
        graph = StartNode >> EndNode

        class Outputs(BaseWorkflow.Outputs):
            final_result = EndNode.Outputs.final

    class TestWorkflowDisplay(BaseWorkflowDisplay[TestWorkflow]):
        edge_displays = {
            (StartNode.Ports.default, EndNode): EdgeDisplay(id=UUID("12345678-1234-5678-1234-567812345678"), z_index=5)
        }

    # WHEN we serialize the workflow with the custom display
    display = get_workflow_display(
        base_display_class=TestWorkflowDisplay,
        workflow_class=TestWorkflow,
    )
    serialized_workflow = display.serialize()

    # THEN the edge should include display_data with z_index
    workflow_raw_data = cast(Dict[str, Any], serialized_workflow["workflow_raw_data"])
    edges = cast(List[Dict[str, Any]], workflow_raw_data["edges"])

    edge_with_display_data = None
    for edge in edges:
        if edge["id"] == "12345678-1234-5678-1234-567812345678":
            edge_with_display_data = edge
            break

    assert edge_with_display_data is not None, "Edge with custom UUID not found"
    assert edge_with_display_data["display_data"] == {"z_index": 5}

    assert edge_with_display_data["type"] == "DEFAULT"
    assert "source_node_id" in edge_with_display_data
    assert "target_node_id" in edge_with_display_data


def test_serialize_workflow_with_node_display_data():
    """
    Tests that nodes with z_index values serialize display_data correctly.
    """

    # GIVEN a workflow with a node that has custom display data
    class TestNode(BaseNode):
        class Outputs(BaseNode.Outputs):
            result: str

    class TestWorkflow(BaseWorkflow):
        graph = TestNode

        class Outputs(BaseWorkflow.Outputs):
            final_result = TestNode.Outputs.result

    class TestNodeDisplay(BaseNodeDisplay[TestNode]):
        display_data = NodeDisplayData(position=NodeDisplayPosition(x=100, y=200), z_index=10, width=300, height=150)

    class TestWorkflowDisplay(BaseWorkflowDisplay[TestWorkflow]):
        pass

    # WHEN we serialize the workflow with the custom node display
    display = get_workflow_display(
        base_display_class=TestWorkflowDisplay,
        workflow_class=TestWorkflow,
    )
    serialized_workflow = display.serialize()

    # THEN the node should include display_data with z_index
    workflow_raw_data = cast(Dict[str, Any], serialized_workflow["workflow_raw_data"])
    nodes = cast(List[Dict[str, Any]], workflow_raw_data["nodes"])

    test_node = None
    for node in nodes:
        if node.get("type") == "GENERIC":
            definition = node.get("definition")
            if isinstance(definition, dict) and definition.get("name") == "TestNode":
                test_node = node
                break

    assert test_node is not None, "TestNode not found in serialized nodes"
    assert test_node["display_data"] == {"position": {"x": 100, "y": 200}, "z_index": 10, "width": 300, "height": 150}


def test_serialize_workflow_with_node_icon_and_color():
    """
    Tests that nodes with icon and color serialize correctly in workflow context.
    """

    # GIVEN a workflow with a node that has icon and color
    class TestNode(BaseNode):
        class Outputs(BaseNode.Outputs):
            result: str

    class TestWorkflow(BaseWorkflow):
        graph = TestNode

        class Outputs(BaseWorkflow.Outputs):
            final_result = TestNode.Outputs.result

    class TestNodeDisplay(BaseNodeDisplay[TestNode]):
        display_data = NodeDisplayData(position=NodeDisplayPosition(x=100, y=200), icon="vellum:icon:cog", color="navy")

    class TestWorkflowDisplay(BaseWorkflowDisplay[TestWorkflow]):
        pass

    # WHEN we serialize the workflow
    display = get_workflow_display(
        base_display_class=TestWorkflowDisplay,
        workflow_class=TestWorkflow,
    )
    serialized_workflow = display.serialize()

    # THEN the node should include icon and color in display_data
    workflow_raw_data = cast(Dict[str, Any], serialized_workflow["workflow_raw_data"])
    nodes = cast(List[Dict[str, Any]], workflow_raw_data["nodes"])

    test_node = None
    for node in nodes:
        if node.get("type") == "GENERIC":
            definition = node.get("definition")
            if isinstance(definition, dict) and definition.get("name") == "TestNode":
                test_node = node
                break

    assert test_node is not None, "TestNode not found in serialized nodes"
    assert test_node["display_data"]["icon"] == "vellum:icon:cog"
    assert test_node["display_data"]["color"] == "navy"


def test_base_workflow_display__graph_with_trigger_and_regular_node():
    """
    Tests that a workflow with both a trigger edge and a regular node edge serializes correctly.
    """

    # GIVEN a workflow with a trigger and regular nodes
    class SlackMessageTrigger(IntegrationTrigger):
        message: str

        class Config(IntegrationTrigger.Config):
            provider = VellumIntegrationProviderType.COMPOSIO
            integration_name = "SLACK"
            slug = "slack_message"

    class TopNode(BaseNode):
        class Outputs(BaseNode.Outputs):
            result: str

    class BottomNode(BaseNode):
        class Outputs(BaseNode.Outputs):
            final: str

    class TestWorkflow(BaseWorkflow):
        graph = {
            Graph.from_node(TopNode),
            SlackMessageTrigger >> BottomNode,
        }

    # WHEN we serialize the workflow
    workflow_display = get_workflow_display(workflow_class=TestWorkflow)
    exec_config = workflow_display.serialize()

    # THEN the serialized workflow should have the expected structure
    workflow_raw_data = exec_config["workflow_raw_data"]
    assert isinstance(workflow_raw_data, dict)

    nodes = workflow_raw_data["nodes"]
    edges = workflow_raw_data["edges"]
    triggers = exec_config["triggers"]

    assert isinstance(nodes, list)
    assert isinstance(edges, list)
    assert isinstance(triggers, list)

    assert len(triggers) == 1
    trigger = triggers[0]
    assert isinstance(trigger, dict)
    trigger_id = trigger["id"]

    entrypoint_nodes = [n for n in nodes if isinstance(n, dict) and n.get("type") == "ENTRYPOINT"]
    assert len(entrypoint_nodes) == 1
    entrypoint_node = entrypoint_nodes[0]
    assert isinstance(entrypoint_node, dict)
    entrypoint_node_id = entrypoint_node["id"]

    top_node = next(
        (
            node
            for node in nodes
            if isinstance(node, dict)
            and isinstance(node.get("definition"), dict)
            and cast(Dict[str, Any], node.get("definition")).get("name") == "TopNode"
        ),
        None,
    )
    bottom_node = next(
        (
            node
            for node in nodes
            if isinstance(node, dict)
            and isinstance(node.get("definition"), dict)
            and cast(Dict[str, Any], node.get("definition")).get("name") == "BottomNode"
        ),
        None,
    )

    assert top_node is not None, "TopNode not found in serialized nodes"
    assert bottom_node is not None, "BottomNode not found in serialized nodes"
    top_node_id = top_node["id"]
    bottom_node_id = bottom_node["id"]

    entrypoint_edges = [e for e in edges if isinstance(e, dict) and e.get("source_node_id") == entrypoint_node_id]
    assert len(entrypoint_edges) == 1, "Should have exactly one edge from the entrypoint node"

    entrypoint_edge = entrypoint_edges[0]
    assert isinstance(entrypoint_edge, dict)
    assert entrypoint_edge["target_node_id"] == top_node_id, "Entrypoint edge should connect to TopNode"

    trigger_edges = [e for e in edges if isinstance(e, dict) and e.get("source_node_id") == trigger_id]
    assert len(trigger_edges) == 1, "Should have exactly one edge from the trigger"

    trigger_edge = trigger_edges[0]
    assert isinstance(trigger_edge, dict)
    assert trigger_edge["target_node_id"] == bottom_node_id, "Trigger edge should connect to BottomNode"


def test_serialize_subworkflow_output_reference_without_display():
    """
    Tests that a workflow with a subworkflow node followed by a generic node
    that references the subworkflow output can be serialized correctly without a display directory.
    """

    # GIVEN a workflow module with a subworkflow node and a follow-on node
    module_path = "tests.workflows.subworkflow_output_reference_without_display"

    # WHEN we serialize the module
    result = BaseWorkflowDisplay.serialize_module(module_path)

    # THEN it should serialize successfully
    assert hasattr(result, "exec_config")
    assert isinstance(result.exec_config, dict)
    assert "workflow_raw_data" in result.exec_config

    # AND the workflow should have the expected nodes
    workflow_raw_data = result.exec_config["workflow_raw_data"]
    assert isinstance(workflow_raw_data, dict)
    nodes = workflow_raw_data.get("nodes")
    assert isinstance(nodes, list)

    subworkflow_node = next(
        (
            node
            for node in nodes
            if isinstance(node, dict)
            and isinstance(node.get("definition"), dict)
            and node.get("definition", {}).get("name") == "SubworkflowNodeExample"
        ),
        None,
    )
    followon_node = next(
        (
            node
            for node in nodes
            if isinstance(node, dict)
            and isinstance(node.get("definition"), dict)
            and node.get("definition", {}).get("name") == "FollowOnNode"
        ),
        None,
    )

    assert subworkflow_node is not None, "SubworkflowNodeExample not found in serialized nodes"
    assert followon_node is not None, "FollowOnNode not found in serialized nodes"

    # AND the subworkflow node should have output variables
    subworkflow_data = subworkflow_node.get("data", {})
    subworkflow_outputs = subworkflow_data.get("output_variables", [])
    assert len(subworkflow_outputs) > 0, "SubworkflowNodeExample should have output variables"

    result_output = next((out for out in subworkflow_outputs if out.get("key") == "result"), None)
    assert result_output is not None, "SubworkflowNodeExample should have 'result' output"
    expected_output_id = result_output["id"]

    # AND the follow-on node should have an attribute that references the subworkflow output
    followon_attributes = followon_node.get("attributes", [])
    assert len(followon_attributes) > 0, "FollowOnNode should have attributes"

    input_value_attr = next(
        (attr for attr in followon_attributes if isinstance(attr, dict) and attr.get("name") == "input_value"), None
    )

    assert input_value_attr is not None, "input_value attribute not found in FollowOnNode"

    # AND the attribute value should reference the subworkflow node's output
    value = input_value_attr.get("value")
    assert isinstance(value, dict)
    assert value.get("type") == "NODE_OUTPUT"
    assert value.get("node_id") == subworkflow_node.get("id")

    # AND the node_output_id should match the actual output ID from the subworkflow
    actual_output_id = value.get("node_output_id")
    assert actual_output_id == expected_output_id, (
        f"node_output_id mismatch: expected {expected_output_id}, got {actual_output_id}. "
        "The reference to the subworkflow output is broken."
    )
