from datetime import datetime

from vellum.client.types.prompt_parameters import PromptParameters
from vellum.client.types.release_review_reviewer import ReleaseReviewReviewer
from vellum.client.types.workflow_deployment_release import (
    ReleaseEnvironment,
    ReleaseReleaseTag,
    SlimReleaseReview,
    WorkflowDeploymentRelease,
    WorkflowDeploymentReleaseWorkflowDeployment,
    WorkflowDeploymentReleaseWorkflowVersion,
)
from vellum.workflows import BaseWorkflow
from vellum.workflows.inputs import BaseInputs
from vellum.workflows.nodes.bases import BaseNode
from vellum.workflows.nodes.displayable.code_execution_node.node import CodeExecutionNode
from vellum.workflows.nodes.displayable.inline_prompt_node.node import InlinePromptNode
from vellum.workflows.nodes.displayable.tool_calling_node.node import ToolCallingNode
from vellum.workflows.nodes.displayable.tool_calling_node.state import ToolCallingState
from vellum.workflows.nodes.displayable.tool_calling_node.utils import create_router_node, create_tool_prompt_node
from vellum.workflows.outputs.base import BaseOutputs
from vellum.workflows.state.base import BaseState
from vellum.workflows.types.definition import (
    AuthorizationType,
    DeploymentDefinition,
    EnvironmentVariableReference,
    MCPServer,
)
from vellum_ee.workflows.display.nodes.get_node_display_class import get_node_display_class
from vellum_ee.workflows.display.workflows.get_vellum_workflow_display_class import get_workflow_display


def test_serialize_node__prompt_inputs__constant_value():
    # GIVEN a prompt node with constant value inputs
    class MyPromptNode(ToolCallingNode):
        prompt_inputs = {"foo": "bar"}

    # AND a workflow with the prompt node
    class Workflow(BaseWorkflow):
        graph = MyPromptNode

    # WHEN the workflow is serialized
    workflow_display = get_workflow_display(workflow_class=Workflow)
    serialized_workflow: dict = workflow_display.serialize()

    # THEN the node should properly serialize the inputs
    my_prompt_node = next(
        node for node in serialized_workflow["workflow_raw_data"]["nodes"] if node["id"] == str(MyPromptNode.__id__)
    )

    prompt_inputs_attribute = next(
        attribute for attribute in my_prompt_node["attributes"] if attribute["name"] == "prompt_inputs"
    )

    assert prompt_inputs_attribute == {
        "id": "fb85d86d-f291-4a0d-b867-f7545df7af59",
        "name": "prompt_inputs",
        "value": {"type": "CONSTANT_VALUE", "value": {"type": "JSON", "value": {"foo": "bar"}}},
    }


def test_serialize_node__prompt_inputs__input_reference():
    # GIVEN a state definition
    class MyInput(BaseInputs):
        foo: str

    # AND a prompt node with inputs
    class MyPromptNode(InlinePromptNode):
        prompt_inputs = {"foo": MyInput.foo}

    # AND a workflow with the prompt node
    class Workflow(BaseWorkflow[MyInput, BaseState]):
        graph = MyPromptNode

    # WHEN the workflow is serialized
    workflow_display = get_workflow_display(workflow_class=Workflow)
    serialized_workflow: dict = workflow_display.serialize()

    # THEN the node should skip the state reference input rule
    my_prompt_node = next(
        node for node in serialized_workflow["workflow_raw_data"]["nodes"] if node["id"] == str(MyPromptNode.__id__)
    )

    prompt_inputs_attribute = next(
        attribute for attribute in my_prompt_node["attributes"] if attribute["name"] == "prompt_inputs"
    )

    assert prompt_inputs_attribute == {
        "id": "80ed13f9-64d2-47ee-bb91-3378de7ad2c0",
        "name": "prompt_inputs",
        "value": {
            "type": "DICTIONARY_REFERENCE",
            "entries": [
                {
                    "id": "981b8cdf-c08d-42a1-a226-76de8acf192f",
                    "key": "foo",
                    "value": {"type": "WORKFLOW_INPUT", "input_variable_id": "e3657390-fd3c-4fea-8cdd-fc5ea79f3278"},
                }
            ],
        },
    }


def test_serialize_node__prompt_inputs__mixed_values():
    # GIVEN a prompt node with mixed values
    class MyInput(BaseInputs):
        foo: str

    # AND a prompt node with mixed values
    class MyPromptNode(InlinePromptNode):
        prompt_inputs = {"foo": "bar", "baz": MyInput.foo}

    # AND a workflow with the prompt node
    class Workflow(BaseWorkflow[MyInput, BaseState]):
        graph = MyPromptNode

    # WHEN the workflow is serialized
    workflow_display = get_workflow_display(workflow_class=Workflow)
    serialized_workflow: dict = workflow_display.serialize()

    # THEN the node should properly serialize the inputs
    my_prompt_node = next(
        node for node in serialized_workflow["workflow_raw_data"]["nodes"] if node["id"] == str(MyPromptNode.__id__)
    )

    prompt_inputs_attribute = next(
        attribute for attribute in my_prompt_node["attributes"] if attribute["name"] == "prompt_inputs"
    )

    assert prompt_inputs_attribute == {
        "id": "7352d310-204c-4291-8757-a84a6e68591a",
        "name": "prompt_inputs",
        "value": {
            "type": "DICTIONARY_REFERENCE",
            "entries": [
                {
                    "id": "05c092c7-4031-43b7-8c3d-b1a317ca271d",
                    "key": "foo",
                    "value": {"type": "CONSTANT_VALUE", "value": {"type": "STRING", "value": "bar"}},
                },
                {
                    "id": "b0de6603-fcdd-44a3-b33a-56f05bd03bb4",
                    "key": "baz",
                    "value": {"type": "WORKFLOW_INPUT", "input_variable_id": "8d57cf1d-147c-427b-9a5e-e5f6ab76e2eb"},
                },
            ],
        },
    }


def test_serialize_node__tool_calling_node__mcp_server_api_key():
    # GIVEN a tool calling node with an mcp server
    class MyToolCallingNode(ToolCallingNode):
        functions = [
            MCPServer(
                name="my-mcp-server",
                url="https://my-mcp-server.com",
                authorization_type=AuthorizationType.API_KEY,
                api_key_header_key="my-api-key-header-key",
                api_key_header_value=EnvironmentVariableReference(name="my-api-key-header-value"),
            )
        ]

    # AND a workflow with the tool calling node
    class Workflow(BaseWorkflow):
        graph = MyToolCallingNode

    # WHEN the workflow is serialized
    workflow_display = get_workflow_display(workflow_class=Workflow)
    serialized_workflow: dict = workflow_display.serialize()

    # THEN the node should properly serialize the mcp server
    my_tool_calling_node = next(
        node
        for node in serialized_workflow["workflow_raw_data"]["nodes"]
        if node["id"] == str(MyToolCallingNode.__id__)
    )

    functions_attribute = next(
        attribute for attribute in my_tool_calling_node["attributes"] if attribute["name"] == "functions"
    )

    assert functions_attribute == {
        "id": "ff00c2d6-f99c-458b-9bcd-181f8e43b2d1",
        "name": "functions",
        "value": {
            "type": "CONSTANT_VALUE",
            "value": {
                "type": "JSON",
                "value": [
                    {
                        "type": "MCP_SERVER",
                        "name": "my-mcp-server",
                        "description": "",
                        "url": "https://my-mcp-server.com",
                        "authorization_type": "API_KEY",
                        "bearer_token_value": None,
                        "api_key_header_key": "my-api-key-header-key",
                        "api_key_header_value": "my-api-key-header-value",
                    }
                ],
            },
        },
    }


def test_serialize_node__tool_calling_node__mcp_server_no_authorization():
    # GIVEN a tool calling node with an mcp server
    class MyToolCallingNode(ToolCallingNode):
        functions = [
            MCPServer(
                name="my-mcp-server",
                url="https://my-mcp-server.com",
            )
        ]

    # AND a workflow with the tool calling node
    class Workflow(BaseWorkflow):
        graph = MyToolCallingNode

    # WHEN the workflow is serialized
    workflow_display = get_workflow_display(workflow_class=Workflow)
    serialized_workflow: dict = workflow_display.serialize()

    # THEN the node should properly serialize the mcp server
    my_tool_calling_node = next(
        node
        for node in serialized_workflow["workflow_raw_data"]["nodes"]
        if node["id"] == str(MyToolCallingNode.__id__)
    )

    functions_attribute = next(
        attribute for attribute in my_tool_calling_node["attributes"] if attribute["name"] == "functions"
    )

    assert functions_attribute == {
        "id": "a8a0133f-7913-451a-89da-f8dea5f352a2",
        "name": "functions",
        "value": {
            "type": "CONSTANT_VALUE",
            "value": {
                "type": "JSON",
                "value": [
                    {
                        "type": "MCP_SERVER",
                        "name": "my-mcp-server",
                        "description": "",
                        "url": "https://my-mcp-server.com",
                        "authorization_type": None,
                        "bearer_token_value": None,
                        "api_key_header_key": None,
                        "api_key_header_value": None,
                    }
                ],
            },
        },
    }


def test_serialize_tool_router_node():
    """
    Test that the tool router node created by create_router_node serializes successfully.
    """

    # GIVEN a simple function for tool calling
    def my_function(arg1: str) -> str:
        return f"Result: {arg1}"

    tool_prompt_node = create_tool_prompt_node(
        ml_model="gpt-4o-mini",
        blocks=[],
        functions=[my_function],
        prompt_inputs=None,
        parameters=PromptParameters(),
    )

    # WHEN we create a router node using create_router_node
    router_node = create_router_node(
        functions=[my_function],
        tool_prompt_node=tool_prompt_node,
    )

    router_node_display_class = get_node_display_class(router_node)
    router_node_display = router_node_display_class()

    class Workflow(BaseWorkflow[BaseInputs, ToolCallingState]):
        graph = tool_prompt_node >> router_node

    workflow_display = get_workflow_display(workflow_class=Workflow)
    display_context = workflow_display.display_context

    # WHEN we serialize the router node
    serialized_router_node = router_node_display.serialize(display_context)

    # THEN the router node should serialize to the exact expected structure
    assert serialized_router_node == {
        "adornments": None,
        "attributes": [
            {
                "id": "ea36d8ff-3f6e-41eb-8a7b-ba3ca8e05f49",
                "name": "prompt_outputs",
                "value": {
                    "node_id": "c75146e1-ea10-4f58-90fd-887322725496",
                    "node_output_id": "baef5c93-612a-453d-b739-223041ef0429",
                    "type": "NODE_OUTPUT",
                },
            }
        ],
        "base": {
            "module": ["vellum", "workflows", "nodes", "displayable", "tool_calling_node", "utils"],
            "name": "RouterNode",
        },
        "definition": {
            "module": ["vellum", "workflows", "nodes", "displayable", "tool_calling_node", "utils"],
            "name": "RouterNode",
        },
        "display_data": {"position": {"x": 0.0, "y": 0.0}},
        "id": "d2884fa0-2d3d-4220-a335-bdcef56a00d5",
        "label": "Router Node",
        "outputs": [],
        "ports": [
            {
                "expression": {
                    "lhs": {
                        "lhs": {
                            "lhs": {
                                "state_variable_id": "0dd7f5a1-1d73-4153-9191-ca828ace4920",
                                "type": "WORKFLOW_STATE",
                            },
                            "operator": "<",
                            "rhs": {
                                "lhs": {
                                    "node_id": "c75146e1-ea10-4f58-90fd-887322725496",
                                    "node_output_id": "baef5c93-612a-453d-b739-223041ef0429",
                                    "type": "NODE_OUTPUT",
                                },
                                "operator": "length",
                                "type": "UNARY_EXPRESSION",
                            },
                            "type": "BINARY_EXPRESSION",
                        },
                        "operator": "and",
                        "rhs": {
                            "lhs": {
                                "lhs": {
                                    "lhs": {
                                        "node_id": "c75146e1-ea10-4f58-90fd-887322725496",
                                        "node_output_id": "baef5c93-612a-453d-b739-223041ef0429",
                                        "type": "NODE_OUTPUT",
                                    },
                                    "operator": "accessField",
                                    "rhs": {
                                        "state_variable_id": "0dd7f5a1-1d73-4153-9191-ca828ace4920",
                                        "type": "WORKFLOW_STATE",
                                    },
                                    "type": "BINARY_EXPRESSION",
                                },
                                "operator": "accessField",
                                "rhs": {"type": "CONSTANT_VALUE", "value": {"type": "STRING", "value": "type"}},
                                "type": "BINARY_EXPRESSION",
                            },
                            "operator": "=",
                            "rhs": {"type": "CONSTANT_VALUE", "value": {"type": "STRING", "value": "FUNCTION_CALL"}},
                            "type": "BINARY_EXPRESSION",
                        },
                        "type": "BINARY_EXPRESSION",
                    },
                    "operator": "and",
                    "rhs": {
                        "lhs": {
                            "lhs": {
                                "lhs": {
                                    "lhs": {
                                        "node_id": "c75146e1-ea10-4f58-90fd-887322725496",
                                        "node_output_id": "baef5c93-612a-453d-b739-223041ef0429",
                                        "type": "NODE_OUTPUT",
                                    },
                                    "operator": "accessField",
                                    "rhs": {
                                        "state_variable_id": "0dd7f5a1-1d73-4153-9191-ca828ace4920",
                                        "type": "WORKFLOW_STATE",
                                    },
                                    "type": "BINARY_EXPRESSION",
                                },
                                "operator": "accessField",
                                "rhs": {"type": "CONSTANT_VALUE", "value": {"type": "STRING", "value": "value"}},
                                "type": "BINARY_EXPRESSION",
                            },
                            "operator": "accessField",
                            "rhs": {"type": "CONSTANT_VALUE", "value": {"type": "STRING", "value": "name"}},
                            "type": "BINARY_EXPRESSION",
                        },
                        "operator": "=",
                        "rhs": {"type": "CONSTANT_VALUE", "value": {"type": "STRING", "value": "my_function"}},
                        "type": "BINARY_EXPRESSION",
                    },
                    "type": "BINARY_EXPRESSION",
                },
                "id": "81d85d92-7f51-4f41-84be-8e2b0eeb4f59",
                "name": "my_function",
                "type": "IF",
            },
            {"expression": None, "id": "a433bba4-a0d6-44f3-b53b-0d5231e442cf", "name": "default", "type": "ELSE"},
        ],
        "trigger": {"id": "9055a5d0-68a1-40cf-bc05-a8c65bd19abe", "merge_behavior": "AWAIT_ATTRIBUTES"},
        "type": "GENERIC",
    }


def test_serialize_node__tool_calling_node__subworkflow_with_parent_input_reference():
    """
    Test that a tool calling node with a subworkflow that references parent inputs serializes correctly
    """

    # GIVEN a workflow inputs class
    class MyInputs(BaseInputs):
        text: str
        greeting: str

    # AND a code execution node that references parent inputs
    class CodeExecution(CodeExecutionNode[BaseState, str]):
        code = ""
        code_inputs = {
            "text": MyInputs.text,
        }
        runtime = "PYTHON_3_11_6"
        packages = []

    # AND a subworkflow that uses the code execution node
    class FunctionSubworkflow(BaseWorkflow):
        graph = CodeExecution

        class Outputs(BaseWorkflow.Outputs):
            output = CodeExecution.Outputs.result

    # AND a tool calling node that uses the subworkflow
    class MyToolCallingNode(ToolCallingNode):
        ml_model = "gpt-4.1"
        prompt_inputs = {
            "text": MyInputs.text,
        }
        blocks = []
        parameters = PromptParameters()
        functions = [FunctionSubworkflow]

    # AND a workflow with the tool calling node
    class Workflow(BaseWorkflow[MyInputs, BaseState]):
        graph = MyToolCallingNode

    # WHEN the workflow is serialized
    workflow_display = get_workflow_display(workflow_class=Workflow)
    serialized_workflow: dict = workflow_display.serialize()

    # THEN the node should properly serialize the functions with parent input references
    my_tool_calling_node = next(
        node
        for node in serialized_workflow["workflow_raw_data"]["nodes"]
        if node["id"] == str(MyToolCallingNode.__id__)
    )

    functions_attribute = next(
        attribute for attribute in my_tool_calling_node["attributes"] if attribute["name"] == "functions"
    )

    data = functions_attribute["value"]["value"]["value"][0]
    nodes = data["exec_config"]["workflow_raw_data"]["nodes"]
    code_exec_node = next((node for node in nodes if node["type"] == "CODE_EXECUTION"), None)

    assert code_exec_node is not None
    text_input = next((input for input in code_exec_node["inputs"] if input["key"] == "text"), None)
    assert text_input == {
        "id": "9f3a85b5-a0c7-4b72-beef-d68b4eb2a428",
        "key": "text",
        "value": {
            "rules": [
                {"type": "INPUT_VARIABLE", "data": {"input_variable_id": "6f0c1889-3f08-4c5c-bb24-f7b94169105c"}}
            ],
            "combinator": "OR",
        },
    }


def test_serialize_tool_prompt_node_with_inline_workflow():
    """
    Test that the tool prompt node created by create_tool_prompt_node serializes successfully with inline workflow.
    """

    # GIVEN a simple inline workflow for tool calling
    class SimpleWorkflowInputs(BaseInputs):
        message: str

    class SimpleNode(BaseNode):
        message = SimpleWorkflowInputs.message

        class Outputs(BaseOutputs):
            result: str

        def run(self) -> Outputs:
            return self.Outputs(result=f"Processed: {self.message}")

    class SimpleInlineWorkflow(BaseWorkflow[SimpleWorkflowInputs, BaseState]):
        """A simple workflow for testing inline tool serialization."""

        graph = SimpleNode

        class Outputs(BaseOutputs):
            result = SimpleNode.Outputs.result

    # WHEN we create a tool prompt node using create_tool_prompt_node with inline workflow
    tool_prompt_node = create_tool_prompt_node(
        ml_model="gpt-4o-mini",
        blocks=[],
        functions=[SimpleInlineWorkflow],
        prompt_inputs=None,
        parameters=PromptParameters(),
    )

    tool_prompt_node_display_class = get_node_display_class(tool_prompt_node)
    tool_prompt_node_display = tool_prompt_node_display_class()

    # AND we create a workflow that uses this tool prompt node
    class TestWorkflow(BaseWorkflow[BaseInputs, ToolCallingState]):
        graph = tool_prompt_node

    # WHEN we serialize the entire workflow
    workflow_display = get_workflow_display(workflow_class=TestWorkflow)
    display_context = workflow_display.display_context
    serialized_tool_prompt_node = tool_prompt_node_display.serialize(display_context)

    # THEN prompt inputs should be serialized correctly
    attributes = serialized_tool_prompt_node["attributes"]
    assert isinstance(attributes, list)
    prompt_inputs_attr = next(
        (attr for attr in attributes if isinstance(attr, dict) and attr["name"] == "prompt_inputs"), None
    )
    assert prompt_inputs_attr == {
        "id": "d218d7c7-41f4-46d4-a9ed-89640cba1b9b",
        "name": "prompt_inputs",
        "value": {
            "type": "DICTIONARY_REFERENCE",
            "entries": [
                {
                    "id": "c181c554-a97f-4254-9bc8-1c30a8f3fb5f",
                    "key": "chat_history",
                    "value": {
                        "type": "BINARY_EXPRESSION",
                        "lhs": {"type": "CONSTANT_VALUE", "value": {"type": "JSON", "value": []}},
                        "operator": "concat",
                        "rhs": {
                            "type": "WORKFLOW_STATE",
                            "state_variable_id": "7a1caaf5-99df-487a-8b2d-6512df2d871a",
                        },
                    },
                }
            ],
        },
    }


def test_serialize_tool_prompt_node_with_workflow_deployment(vellum_client):
    """
    Test that the tool prompt node serializes successfully with a workflow deployment.
    """
    vellum_client.workflow_deployments.retrieve_workflow_deployment_release.return_value = WorkflowDeploymentRelease(
        id="test-id",
        created=datetime.now(),
        environment=ReleaseEnvironment(
            id="test-id",
            name="test-name",
            label="test-label",
        ),
        created_by=None,
        workflow_version=WorkflowDeploymentReleaseWorkflowVersion(
            id="test-id",
            input_variables=[],
            output_variables=[],
        ),
        deployment=WorkflowDeploymentReleaseWorkflowDeployment(name="test-name"),
        description="test-description",
        release_tags=[
            ReleaseReleaseTag(
                name="test-name",
                source="USER",
            )
        ],
        reviews=[
            SlimReleaseReview(
                id="test-id",
                created=datetime.now(),
                reviewer=ReleaseReviewReviewer(
                    id="test-id",
                    full_name="test-name",
                ),
                state="APPROVED",
            )
        ],
    )

    # GIVEN a workflow deployment
    workflow_deployment = DeploymentDefinition(
        deployment="test-deployment",
        release_tag="test-release-tag",
    )

    # WHEN we create a tool prompt node using create_tool_prompt_node with a workflow deployment
    tool_prompt_node = create_tool_prompt_node(
        ml_model="gpt-4o-mini",
        blocks=[],
        functions=[workflow_deployment],
        prompt_inputs=None,
        parameters=PromptParameters(),
    )

    tool_prompt_node_display_class = get_node_display_class(tool_prompt_node)
    tool_prompt_node_display = tool_prompt_node_display_class()

    # AND we create a workflow that uses this tool prompt node
    class TestWorkflow(BaseWorkflow[BaseInputs, ToolCallingState]):
        graph = tool_prompt_node

    # WHEN we serialize the entire workflow
    workflow_display = get_workflow_display(workflow_class=TestWorkflow)
    display_context = workflow_display.display_context
    serialized_tool_prompt_node = tool_prompt_node_display.serialize(display_context)

    # THEN functions attribute should be serialized correctly
    attributes = serialized_tool_prompt_node["attributes"]
    assert isinstance(attributes, list)
    functions_attr = next((attr for attr in attributes if isinstance(attr, dict) and attr["name"] == "functions"), None)
    assert functions_attr == {
        "id": "f482dc81-e320-402d-83df-f60d278797d8",
        "name": "functions",
        "value": {
            "type": "CONSTANT_VALUE",
            "value": {
                "type": "JSON",
                "value": [
                    {
                        "type": "WORKFLOW_DEPLOYMENT",
                        "name": "test-name",
                        "description": "test-description",
                        "deployment": "test-deployment",
                        "release_tag": "test-release-tag",
                    }
                ],
            },
        },
    }
