import pytest
import os
import re
from typing import Any, List, Union

from pydantic import BaseModel

from vellum import ArrayInput, CodeExecutorResponse, NumberVellumValue, StringInput, StringVellumValue
from vellum.client.core.api_error import ApiError
from vellum.client.errors.bad_request_error import BadRequestError
from vellum.client.errors.forbidden_error import ForbiddenError
from vellum.client.errors.internal_server_error import InternalServerError
from vellum.client.errors.not_found_error import NotFoundError
from vellum.client.types.chat_message import ChatMessage
from vellum.client.types.code_execution_package import CodeExecutionPackage
from vellum.client.types.code_executor_secret_input import CodeExecutorSecretInput
from vellum.client.types.function_call import FunctionCall
from vellum.client.types.number_input import NumberInput
from vellum.client.types.string_chat_message_content import StringChatMessageContent
from vellum.workflows.constants import undefined
from vellum.workflows.errors import WorkflowErrorCode
from vellum.workflows.exceptions import NodeException
from vellum.workflows.inputs.base import BaseInputs
from vellum.workflows.nodes.displayable.code_execution_node import CodeExecutionNode
from vellum.workflows.references.vellum_secret import VellumSecretReference
from vellum.workflows.state.base import BaseState, StateMeta
from vellum.workflows.types.core import Json


def test_run_node__happy_path(vellum_client):
    """Confirm that CodeExecutionNodes output the expected text and results when run."""

    # GIVEN a node that subclasses CodeExecutionNode
    class Inputs(BaseInputs):
        word: str

    class State(BaseState):
        pass

    fixture = os.path.abspath(os.path.join(__file__, "../fixtures/main.py"))

    class ExampleCodeExecutionNode(CodeExecutionNode[State, int]):
        filepath = fixture
        runtime = "PYTHON_3_11_6"
        packages = [
            CodeExecutionPackage(
                name="openai",
                version="1.0.0",
            )
        ]

        code_inputs = {
            "word": Inputs.word,
        }

    # AND we know what the Code Execution Node will respond with
    mock_code_execution = CodeExecutorResponse(
        log="hello",
        output=NumberVellumValue(value=5),
    )
    vellum_client.execute_code.return_value = mock_code_execution

    # WHEN we run the node
    node = ExampleCodeExecutionNode(
        state=State(
            meta=StateMeta(workflow_inputs=Inputs(word="hello")),
        )
    )
    outputs = node.run()

    # THEN the node should have produced the outputs we expect
    assert outputs == {"result": 5, "log": "hello"}

    # AND we should have invoked the Code with the expected inputs
    vellum_client.execute_code.assert_called_once_with(
        input_values=[
            StringInput(name="word", value="hello"),
        ],
        code="""\
def main(word: str) -> int:
    print(word)  # noqa: T201
    return len(word)
""",
        runtime="PYTHON_3_11_6",
        output_type="NUMBER",
        packages=[
            CodeExecutionPackage(
                name="openai",
                version="1.0.0",
            )
        ],
        request_options=None,
    )


def test_run_node__code_attribute(vellum_client):
    """Confirm that CodeExecutionNodes can use the `code` attribute to specify the code to execute."""

    # GIVEN a node that subclasses CodeExecutionNode
    class Inputs(BaseInputs):
        word: str

    class State(BaseState):
        pass

    class ExampleCodeExecutionNode(CodeExecutionNode[State, int]):
        code = """\
def main(word: str) -> int:
    print(word)  # noqa: T201
    return len(word)
"""
        runtime = "PYTHON_3_11_6"
        packages = [
            CodeExecutionPackage(
                name="openai",
                version="1.0.0",
                repository="test-repo",
            )
        ]

        code_inputs = {
            "word": Inputs.word,
        }

    # AND we know what the Code Execution Node will respond with
    mock_code_execution = CodeExecutorResponse(
        log="hello",
        output=NumberVellumValue(value=5),
    )
    vellum_client.execute_code.return_value = mock_code_execution

    # WHEN we run the node
    node = ExampleCodeExecutionNode(
        state=State(
            meta=StateMeta(workflow_inputs=Inputs(word="hello")),
        )
    )
    outputs = node.run()

    # THEN the node should have produced the outputs we expect
    assert outputs == {"result": 5, "log": "hello"}

    # AND we should have invoked the Code with the expected inputs
    vellum_client.execute_code.assert_called_once_with(
        input_values=[
            StringInput(name="word", value="hello"),
        ],
        code="""\
def main(word: str) -> int:
    print(word)  # noqa: T201
    return len(word)
""",
        runtime="PYTHON_3_11_6",
        output_type="NUMBER",
        packages=[
            CodeExecutionPackage(
                name="openai",
                version="1.0.0",
                repository="test-repo",
            )
        ],
        request_options=None,
    )


def test_run_node__code_and_filepath_defined(vellum_client):
    """Confirm that CodeExecutionNodes raise an error if both `code` and `filepath` are defined."""

    # GIVEN a node that subclasses CodeExecutionNode
    class Inputs(BaseInputs):
        word: str

    class State(BaseState):
        pass

    fixture = os.path.abspath(os.path.join(__file__, "../fixtures/main.py"))

    class ExampleCodeExecutionNode(CodeExecutionNode[State, int]):
        filepath = fixture
        code = """\
def main(word: str) -> int:
    print(word)  # noqa: T201
    return len(word)
"""
        runtime = "PYTHON_3_11_6"
        packages = [
            CodeExecutionPackage(
                name="openai",
                version="1.0.0",
            )
        ]

        code_inputs = {
            "word": Inputs.word,
        }

    # AND we know what the Code Execution Node will respond with
    mock_code_execution = CodeExecutorResponse(
        log="hello",
        output=NumberVellumValue(value=5),
    )
    vellum_client.execute_code.return_value = mock_code_execution

    # WHEN we run the node
    node = ExampleCodeExecutionNode(
        state=State(
            meta=StateMeta(workflow_inputs=Inputs(word="hello")),
        )
    )
    with pytest.raises(NodeException) as exc_info:
        node.run()

    # THEN the node should have produced the exception we expected
    assert exc_info.value.message == "Cannot specify both `code` and `filepath` for a CodeExecutionNode"


def test_run_node__code_and_filepath_not_defined(vellum_client):
    """Confirm that CodeExecutionNodes raise an error if neither `code` nor `filepath` are defined."""

    # GIVEN a node that subclasses CodeExecutionNode
    class Inputs(BaseInputs):
        word: str

    class State(BaseState):
        pass

    class ExampleCodeExecutionNode(CodeExecutionNode[State, int]):
        runtime = "PYTHON_3_11_6"
        packages = [
            CodeExecutionPackage(
                name="openai",
                version="1.0.0",
            )
        ]

        code_inputs = {
            "word": Inputs.word,
        }

    # AND we know what the Code Execution Node will respond with
    mock_code_execution = CodeExecutorResponse(
        log="hello",
        output=NumberVellumValue(value=5),
    )
    vellum_client.execute_code.return_value = mock_code_execution

    # WHEN we run the node
    node = ExampleCodeExecutionNode(
        state=State(
            meta=StateMeta(workflow_inputs=Inputs(word="hello")),
        )
    )
    with pytest.raises(NodeException) as exc_info:
        node.run()

    # THEN the node should have produced the exception we expected
    assert exc_info.value.message == "Must specify either `code` or `filepath` for a CodeExecutionNode"


def test_run_node__vellum_secret(vellum_client):
    """Confirm that CodeExecutionNodes can use Vellum Secrets"""

    # GIVEN a node that subclasses CodeExecutionNode that references a Vellum Secret
    class State(BaseState):
        pass

    fixture = os.path.abspath(os.path.join(__file__, "../fixtures/main.py"))

    class ExampleCodeExecutionNode(CodeExecutionNode[State, int]):
        filepath = fixture
        runtime = "PYTHON_3_11_6"
        packages = [
            CodeExecutionPackage(
                name="openai",
                version="1.0.0",
            )
        ]

        code_inputs = {
            "word": VellumSecretReference("OPENAI_API_KEY"),
        }

    # AND we know what the Code Execution Node will respond with
    mock_code_execution = CodeExecutorResponse(
        log="",
        output=NumberVellumValue(value=0),
    )
    vellum_client.execute_code.return_value = mock_code_execution

    # WHEN we run the node
    node = ExampleCodeExecutionNode(state=State())
    outputs = node.run()

    # THEN the node should have produced the outputs we expect
    assert outputs == {"result": 0, "log": ""}

    # AND we should have invoked the Code with the expected inputs
    vellum_client.execute_code.assert_called_once_with(
        input_values=[
            CodeExecutorSecretInput(
                name="word",
                value="OPENAI_API_KEY",
            )
        ],
        code="""\
def main(word: str) -> int:
    print(word)  # noqa: T201
    return len(word)
""",
        runtime="PYTHON_3_11_6",
        output_type="NUMBER",
        packages=[
            CodeExecutionPackage(
                name="openai",
                version="1.0.0",
            )
        ],
        request_options=None,
    )


def test_run_node__int_input(vellum_client):
    """Confirm that CodeExecutionNodes can use int's as inputs"""

    # GIVEN a node that subclasses CodeExecutionNode that references an int
    class State(BaseState):
        pass

    fixture = os.path.abspath(os.path.join(__file__, "../fixtures/main.py"))

    class ExampleCodeExecutionNode(CodeExecutionNode[State, int]):
        filepath = fixture
        runtime = "PYTHON_3_11_6"
        packages = [
            CodeExecutionPackage(
                name="openai",
                version="1.0.0",
            )
        ]

        code_inputs = {
            "counter": 1,
        }

    # AND we know what the Code Execution Node will respond with
    mock_code_execution = CodeExecutorResponse(
        log="",
        output=NumberVellumValue(value=0),
    )
    vellum_client.execute_code.return_value = mock_code_execution

    # WHEN we run the node
    node = ExampleCodeExecutionNode(state=State())
    outputs = node.run()

    # THEN the node should have produced the outputs we expect
    assert outputs == {"result": 0, "log": ""}

    # AND we should have invoked the Code with the correct inputs
    assert vellum_client.execute_code.call_args_list[0].kwargs["input_values"] == [
        NumberInput(
            name="counter",
            value=1.0,
        )
    ]


def test_run_node__run_inline(vellum_client):
    """Confirm that CodeExecutionNodes run the code inline instead of through Vellum under certain conditions."""

    # GIVEN a node that subclasses CodeExecutionNode
    class ExampleCodeExecutionNode(CodeExecutionNode[BaseState, int]):
        code = """\
def main(word: str) -> int:
    print(word)  # noqa: T201
    return len(word)
"""
        runtime = "PYTHON_3_11_6"

        code_inputs = {
            "word": "hello",
        }

    # WHEN we run the node
    node = ExampleCodeExecutionNode()
    outputs = node.run()

    # THEN the node should have produced the outputs we expect
    assert outputs == {"result": 5, "log": "hello\n"}

    # AND we should have not invoked the Code via Vellum
    vellum_client.execute_code.assert_not_called()


def test_run_node__run_inline__incorrect_output_type():
    """Confirm that CodeExecutionNodes raise an error if the output type is incorrect during inline execution."""

    # GIVEN a node that subclasses CodeExecutionNode that returns a string but is defined to return an int
    class ExampleCodeExecutionNode(CodeExecutionNode[BaseState, int]):
        code = """\
def main(word: str) -> int:
    return word
"""
        runtime = "PYTHON_3_11_6"

        code_inputs = {
            "word": "hello",
        }

    # WHEN we run the node
    node = ExampleCodeExecutionNode()
    with pytest.raises(NodeException) as exc_info:
        node.run()

    # THEN the node should have produced the exception we expected
    assert exc_info.value.message == "Expected an output of type 'int', but received 'str'"


def test_run_node__run_inline__incorrect_output_type_list():
    """Confirm that CodeExecutionNodes raise an error if the output type is incorrect during inline execution."""

    # GIVEN a node that subclasses CodeExecutionNode that returns a list but is defined to return an int
    class ExampleCodeExecutionNode(CodeExecutionNode[BaseState, int]):
        code = """\
def main(output: list[str]) -> int:
    return output.value
"""
        runtime = "PYTHON_3_11_6"

        code_inputs = {
            "output": ["hello", "world"],
        }

    # WHEN we run the node
    node = ExampleCodeExecutionNode()
    with pytest.raises(NodeException) as exc_info:
        node.run()

    # THEN the node should have produced the exception we expected
    assert exc_info.value.message == "Expected an output of type 'int', but received 'list'"


def test_run_node__run_inline__incorrect_output_type_dict():
    """Confirm that CodeExecutionNodes raise an error if the output type is incorrect during inline execution."""

    # GIVEN a node that subclasses CodeExecutionNode that returns a dict but is defined to return an int
    class ExampleCodeExecutionNode(CodeExecutionNode[BaseState, int]):
        code = """\
def main(output: dict[str, str]) -> int:
    return output.value
"""
        runtime = "PYTHON_3_11_6"

        code_inputs = {
            "output": {
                "foo": "bar",
            }
        }

    # WHEN we run the node
    node = ExampleCodeExecutionNode()
    with pytest.raises(NodeException) as exc_info:
        node.run()

    # THEN the node should have produced the exception we expected
    assert exc_info.value.message == "Expected an output of type 'int', but received 'dict'"


def test_run_node__run_inline__valid_dict_to_pydantic():
    """Confirm that CodeExecutionNodes can convert a dict to a Pydantic model during inline execution."""

    # GIVEN a node that subclasses CodeExecutionNode that returns a dict matching a Pydantic model
    class ExampleCodeExecutionNode(CodeExecutionNode[BaseState, FunctionCall]):
        code = """\
def main(word: str) -> int:
    return {
        "name": word,
        "arguments": {},
    }
"""
        runtime = "PYTHON_3_11_6"

        code_inputs = {
            "word": "hello",
        }

    # WHEN we run the node
    node = ExampleCodeExecutionNode()
    outputs = node.run()

    # THEN the node should have produced the outputs we expect
    assert outputs == {"result": FunctionCall(name="hello", arguments={}), "log": ""}


def test_run_node__run_inline__invalid_dict_to_pydantic():
    """Confirm that CodeExecutionNodes raise an error if the Pydantic validation fails during inline execution."""

    # GIVEN a node that subclasses CodeExecutionNode that returns a dict not matching a Pydantic model
    class ExampleCodeExecutionNode(CodeExecutionNode[BaseState, FunctionCall]):
        code = """\
def main(word: str) -> int:
    return {
        "n": word,
        "a": {},
    }
"""
        runtime = "PYTHON_3_11_6"

        code_inputs = {
            "word": "hello",
        }

    # WHEN we run the node
    node = ExampleCodeExecutionNode()
    with pytest.raises(NodeException) as exc_info:
        node.run()

    # THEN the node should have produced the exception we expected
    assert exc_info.value.message == "Expected an output of type 'FunctionCall', but received 'dict'"


def test_run_node__run_inline__valid_dict_to_pydantic_any_type():
    """Confirm that CodeExecutionNodes can convert a dict to a Pydantic model during inline execution."""

    # GIVEN a node that subclasses CodeExecutionNode that returns a dict matching Any
    class ExampleCodeExecutionNode(CodeExecutionNode[BaseState, Any]):
        code = """\
def main(word: str) -> dict:
    return {
        "name": "word",
        "arguments": {},
    }
    """
        runtime = "PYTHON_3_11_6"

        code_inputs = {
            "word": "hello",
        }

    # WHEN we run the node
    node = ExampleCodeExecutionNode()
    outputs = node.run()

    # THEN the node should have produced the outputs we expect
    assert outputs == {
        "result": {
            "name": "word",
            "arguments": {},
        },
        "log": "",
    }


def test_run_node__array_input_with_vellum_values(vellum_client):
    """Confirm that CodeExecutionNodes can handle arrays containing VellumValue objects."""

    # GIVEN a node that subclasses CodeExecutionNode that processes an array of VellumValues
    class State(BaseState):
        pass

    class ExampleCodeExecutionNode(CodeExecutionNode[State, str]):
        code = """\
from typing import List, Dict
def main(arg1: List[Dict]) -> str:
    return arg1[0]["value"] + " " + arg1[1]["value"]
"""
        runtime = "PYTHON_3_11_6"

        code_inputs = {
            "arg1": [
                StringVellumValue(type="STRING", value="Hello", name="First"),
                StringVellumValue(type="STRING", value="World", name="Second"),
            ],
        }

    # WHEN we run the node
    node = ExampleCodeExecutionNode(state=State())
    outputs = node.run()

    # THEN the node should successfully concatenate the values
    assert outputs == {"result": "Hello World", "log": ""}

    # AND we should not have invoked the Code via Vellum since it's running inline
    vellum_client.execute_code.assert_not_called()


def test_run_node__union_output_type(vellum_client):
    """Confirm that CodeExecutionNodes can handle Union output types."""

    # GIVEN a node that subclasses CodeExecutionNode that returns a Union type
    class State(BaseState):
        pass

    class ExampleCodeExecutionNode(CodeExecutionNode[State, Union[float, int]]):
        code = """\
from typing import List, Dict
def main(arg1: List[Dict]) -> float:
    return arg1[0]["value"] + arg1[1]["value"]
"""
        runtime = "PYTHON_3_11_6"

        code_inputs = {
            "arg1": [
                NumberVellumValue(type="NUMBER", value=1.0, name="First"),
                NumberVellumValue(type="NUMBER", value=2.0, name="Second"),
            ],
        }

    # WHEN we run the node
    node = ExampleCodeExecutionNode(state=State())
    outputs = node.run()

    # THEN the node should successfully sum the values
    assert outputs == {"result": 3.0, "log": ""}

    # AND we should not have invoked the Code via Vellum since it's running inline
    vellum_client.execute_code.assert_not_called()


def test_run_node__code_execution_error():
    # GIVEN a node that will raise an error during execution
    class State(BaseState):
        pass

    class ExampleCodeExecutionNode(CodeExecutionNode[State, int]):
        code = """\
def main(arg1: int, arg2: int) -> int:
    return arg1 + arg2 + arg3
"""
        runtime = "PYTHON_3_11_6"
        code_inputs = {
            "arg1": 1,
            "arg2": 2,
        }

    # WHEN we run the node
    node = ExampleCodeExecutionNode(state=State())

    # THEN it should raise a NodeException with the execution error
    with pytest.raises(NodeException) as exc_info:
        node.run()

    # AND the error should contain the execution error details
    assert "name 'arg3' is not defined" in str(exc_info.value)
    assert exc_info.value.code == WorkflowErrorCode.INVALID_CODE


def test_run_node__array_of_bools_input():
    # GIVEN a node that will raise an error during execution
    class ExampleCodeExecutionNode(CodeExecutionNode[BaseState, int]):
        code = """\
def main(arg1: list[bool]) -> int:
    return len(arg1)
"""
        runtime = "PYTHON_3_11_6"
        code_inputs = {
            "arg1": [True, False, True],
        }

    # WHEN we run the node
    node = ExampleCodeExecutionNode()

    # THEN it should raise a NodeException with the execution error
    outputs = node.run()

    # AND the error should contain the execution error details
    assert outputs == {"result": 3, "log": ""}


def test_run_node__union_output_type__pydantic_children():
    # GIVEN a node that is a union type with a pydantic child
    class OptionOne(BaseModel):
        foo: str

    class OptionTwo(BaseModel):
        bar: int

    class ExampleCodeExecutionNode(CodeExecutionNode[BaseState, Union[OptionOne, OptionTwo]]):
        code = """\
def main():
    return { "foo": "hello" }
"""
        runtime = "PYTHON_3_11_6"
        code_inputs = {}

    # WHEN we run the node
    node = ExampleCodeExecutionNode()

    # THEN it should run successfully
    outputs = node.run()

    # AND the result should be the correct type
    assert outputs == {"result": OptionOne(foo="hello"), "log": ""}


def test_run_node__union_output_type__miss():
    # GIVEN a node that is a union type
    class ExampleCodeExecutionNode(CodeExecutionNode[BaseState, Union[int, float]]):
        code = """\
def main():
    return "hello"
"""
        runtime = "PYTHON_3_11_6"
        code_inputs = {}

    # WHEN we run the node
    node = ExampleCodeExecutionNode()

    # THEN it should raise a NodeException with the execution error
    with pytest.raises(NodeException) as exc_info:
        node.run()

    # AND the error should contain the execution error details
    assert re.match(
        r"Expected an output of type '(int \| float|float \| int)', but received 'str'", exc_info.value.message
    )


def test_run_node__chat_history_output_type():
    # GIVEN a node that that has a chat history return type
    class ExampleCodeExecutionNode(CodeExecutionNode[BaseState, List[ChatMessage]]):
        code = """\
def main():
    return [
        {
            "role": "USER",
            "content": {
                "type": "STRING",
                "value": "Hello, world!",
            }
        }
    ]
"""
        code_inputs = {}
        runtime = "PYTHON_3_11_6"

    # WHEN we run the node
    node = ExampleCodeExecutionNode()
    outputs = node.run()

    # AND the error should contain the execution error details
    assert outputs == {
        "result": [ChatMessage(role="USER", content=StringChatMessageContent(value="Hello, world!"))],
        "log": "",
    }


def test_run_node__execute_code_api_fails__node_exception(vellum_client):
    # GIVEN a node that will throw a JSON.parse error
    class ExampleCodeExecutionNode(CodeExecutionNode[BaseState, Json]):
        code = """\
async function main(inputs: {
  data: string,
}): Promise<string> {
  return JSON.parse(inputs.data)
}
"""
        code_inputs = {
            "data": "not a valid json string",
        }
        runtime = "TYPESCRIPT_5_3_3"

    # AND the execute_code API will fail
    message = """\
Code execution error (exit code 1): undefined:1
not a valid json string
 ^

SyntaxError: Unexpected token 'o', \"not a valid\"... is not valid JSON
    at JSON.parse (<anonymous>)
    at Object.eval (eval at execute (/workdir/runner.js:16:18), <anonymous>:40:40)
    at step (eval at execute (/workdir/runner.js:16:18), <anonymous>:32:23)
    at Object.eval [as next] (eval at execute (/workdir/runner.js:16:18), <anonymous>:13:53)
    at eval (eval at execute (/workdir/runner.js:16:18), <anonymous>:7:71)
    at new Promise (<anonymous>)
    at __awaiter (eval at execute (/workdir/runner.js:16:18), <anonymous>:3:12)
    at Object.main (eval at execute (/workdir/runner.js:16:18), <anonymous>:38:12)
    at execute (/workdir/runner.js:17:33)
    at Interface.<anonymous> (/workdir/runner.js:58:5)

Node.js v21.7.3
"""
    vellum_client.execute_code.side_effect = BadRequestError(
        body={
            "message": message,
            "log": "",
        }
    )

    # WHEN we run the node
    node = ExampleCodeExecutionNode()
    with pytest.raises(NodeException) as exc_info:
        node.run()

    # AND the error should contain the execution error details
    assert exc_info.value.message == message


def test_run_node__execute_code_api_fails_403__provider_credentials_unavailable(vellum_client):
    """
    Tests that a 403 error from the API is handled with PROVIDER_CREDENTIALS_UNAVAILABLE error code.
    """

    class ExampleCodeExecutionNode(CodeExecutionNode[BaseState, str]):
        code = "def main(): return 'test'"
        runtime = "PYTHON_3_11_6"
        packages = [CodeExecutionPackage(name="requests", version="2.28.0")]

    vellum_client.execute_code.side_effect = ForbiddenError(
        body={
            "detail": "Provider credentials is missing or unavailable",
        }
    )

    node = ExampleCodeExecutionNode()
    with pytest.raises(NodeException) as exc_info:
        node.run()

    assert exc_info.value.code == WorkflowErrorCode.PROVIDER_CREDENTIALS_UNAVAILABLE
    assert "Provider credentials is missing or unavailable" in exc_info.value.message


def test_run_node__execute_code_api_fails_404__invalid_inputs(vellum_client):
    """
    Tests that a 404 error from the API is handled with INVALID_INPUTS error code.
    """

    class ExampleCodeExecutionNode(CodeExecutionNode[BaseState, str]):
        code = "def main(): return 'test'"
        runtime = "PYTHON_3_11_6"
        packages = [CodeExecutionPackage(name="requests", version="2.28.0")]

    vellum_client.execute_code.side_effect = NotFoundError(
        body={
            "detail": "Resource not found",
        }
    )

    node = ExampleCodeExecutionNode()
    with pytest.raises(NodeException) as exc_info:
        node.run()

    assert exc_info.value.code == WorkflowErrorCode.INVALID_INPUTS
    assert "Resource not found" in exc_info.value.message


def test_run_node__execute_code_api_fails_500__internal_error(vellum_client):
    """
    Tests that a 500 error from the API is handled with INTERNAL_ERROR error code.
    """

    class ExampleCodeExecutionNode(CodeExecutionNode[BaseState, str]):
        code = "def main(): return 'test'"
        runtime = "PYTHON_3_11_6"
        packages = [CodeExecutionPackage(name="requests", version="2.28.0")]

    vellum_client.execute_code.side_effect = InternalServerError(
        body={
            "detail": "Internal server error occurred",
        }
    )

    node = ExampleCodeExecutionNode()
    with pytest.raises(NodeException) as exc_info:
        node.run()

    assert exc_info.value.code == WorkflowErrorCode.INTERNAL_ERROR
    assert exc_info.value.message == "Failed to execute code"


def test_run_node__execute_code_api_fails_4xx_no_message__uses_detail(vellum_client):
    """
    Tests that a 4xx error without a 'message' field falls back to 'detail' field.
    """

    class ExampleCodeExecutionNode(CodeExecutionNode[BaseState, str]):
        code = "def main(): return 'test'"
        runtime = "PYTHON_3_11_6"
        packages = [CodeExecutionPackage(name="requests", version="2.28.0")]

    vellum_client.execute_code.side_effect = ApiError(
        status_code=422,
        body={
            "detail": "Invalid request parameters",
        },
    )

    node = ExampleCodeExecutionNode()
    with pytest.raises(NodeException) as exc_info:
        node.run()

    assert exc_info.value.code == WorkflowErrorCode.INVALID_INPUTS
    assert "Invalid request parameters" in exc_info.value.message


def test_run_node__execute_code__list_extends():
    # GIVEN a node that will return a list with output type Json
    class ExampleCodeExecutionNode(CodeExecutionNode[BaseState, Json]):
        code = """\
def main(left, right):
    all = []
    all.extend(left)
    all.extend(right)
    return all
"""
        code_inputs = {
            "left": [1, 2, 3],
            "right": [4, 5, 6],
        }
        runtime = "PYTHON_3_11_6"

    # WHEN we run the node
    node = ExampleCodeExecutionNode()
    outputs = node.run()

    # AND the result should be the correct output
    assert outputs == {"result": [1, 2, 3, 4, 5, 6], "log": ""}


def test_run_node__execute_code__non_str_print():
    # GIVEN a node that will print a non-string value
    class ExampleCodeExecutionNode(CodeExecutionNode[BaseState, str]):
        code = """\
def main():
    print(type(1))
    return "hello"
"""
        code_inputs = {}
        runtime = "PYTHON_3_11_6"

    # WHEN we run the node
    node = ExampleCodeExecutionNode()
    outputs = node.run()

    # AND the result should be the correct output
    assert outputs == {"result": "hello", "log": "<class 'int'>\n"}


def test_run_node__execute_code__invalid_key_access():
    # GIVEN a node that will access an invalid key
    class ExampleCodeExecutionNode(CodeExecutionNode[BaseState, str]):
        code = """\
def main(arg1: list) -> str:
    return arg1["invalid"]
"""
        code_inputs = {
            "arg1": {"foo": "bar"},
        }
        runtime = "PYTHON_3_11_6"

    # WHEN we run the node
    node = ExampleCodeExecutionNode()
    outputs = node.run()

    # AND the result should be the correct output
    assert outputs == {"result": "", "log": ""}


def test_run_node__execute_code__value_key_access():
    # GIVEN a node that will access the value key before the nested key
    class ExampleCodeExecutionNode(CodeExecutionNode[BaseState, str]):
        code = """\
def main(arg1: list) -> str:
    return arg1["value"]["foo"]
"""
        code_inputs = {
            "arg1": {"foo": "bar"},
        }
        runtime = "PYTHON_3_11_6"

    # WHEN we run the node
    node = ExampleCodeExecutionNode()
    outputs = node.run()

    # AND the result should be the correct output
    assert outputs == {"result": "bar", "log": ""}


@pytest.mark.parametrize(
    "code_snippet",
    [
        """
def main(text: str) -> str:
    return text.value
""",
        """
def main(text: str) -> str:
    return text["value"]
""",
    ],
)
def test_run_node__string_value_wrapper_value(code_snippet):
    """Test string value wrapper value access using different patterns"""

    # GIVEN a node that accesses the 'value' property of a string input
    class ExampleCodeExecutionNode(CodeExecutionNode[BaseState, str]):
        code = code_snippet
        runtime = "PYTHON_3_11_6"
        code_inputs = {
            "text": "hello",
        }

    # WHEN we run the node
    node = ExampleCodeExecutionNode()
    outputs = node.run()

    # THEN the node should successfully access the string value
    assert outputs == {"result": "hello", "log": ""}


@pytest.mark.parametrize(
    "code_snippet",
    [
        """
def main(text: str) -> str:
    return text.type
""",
        """
def main(text: str) -> str:
    return text["type"]
""",
    ],
)
def test_run_node__string_value_wrapper_type(code_snippet):
    """Test string value wrapper type access using different patterns"""

    # GIVEN a node that will return the string type
    class ExampleCodeExecutionNode(CodeExecutionNode[BaseState, str]):
        code = code_snippet
        runtime = "PYTHON_3_11_6"
        code_inputs = {
            "text": "hello",
        }

    # WHEN we run the node
    node = ExampleCodeExecutionNode()
    outputs = node.run()

    # THEN the node should successfully return the string type
    assert outputs == {"result": "STRING", "log": ""}


def test_run_node__string_value_wrapper__list_of_dicts():
    # GIVEN a node that accesses the 'value' property of a string input
    class ExampleCodeExecutionNode(CodeExecutionNode[BaseState, Any]):
        code = """\
def main(output: list[str]) -> list[str]:
    results = []
    for item in output:
        results.append(item['value'])

    return results
"""
        code_inputs = {"output": ['{"foo": "bar"}', '{"foo2": "bar2"}']}
        runtime = "PYTHON_3_11_6"

    # WHEN we run the node
    node = ExampleCodeExecutionNode()
    outputs = node.run()

    # THEN the node should successfully access the string value
    assert outputs == {"result": ['{"foo": "bar"}', '{"foo2": "bar2"}'], "log": ""}


def test_run_node__string_key_access_still_works():
    # GIVEN a node that accesses the '0' index of a string input
    class ExampleCodeExecutionNode(CodeExecutionNode[BaseState, Any]):
        code = """\
def main(input: str) -> str:
    return input[0]
"""
        code_inputs = {"input": "hello"}
        runtime = "PYTHON_3_11_6"

    # WHEN we run the node
    node = ExampleCodeExecutionNode()
    outputs = node.run()

    # THEN the node should successfully access the string value
    assert outputs == {"result": "h", "log": ""}


@pytest.mark.parametrize(
    "code_snippet",
    [
        """
from vellum.client.types.function_call import FunctionCall
def main(input: FunctionCall) -> FunctionCall:
    return input.value
""",
        """
from vellum.client.types.function_call import FunctionCall
def main(input: FunctionCall) -> FunctionCall:
    return input["value"]
""",
    ],
)
def test_run_node__function_call_wrapper_value(code_snippet):
    """Test function call wrapper value access using different patterns"""

    # GIVEN a node that accesses the function call value
    class ExampleCodeExecutionNode(CodeExecutionNode[BaseState, FunctionCall]):
        code = code_snippet
        code_inputs = {
            "input": FunctionCall(
                name="test-name",
                arguments={
                    "test-key": "test-value",
                },
            )
        }

    # WHEN we run the node
    node = ExampleCodeExecutionNode()
    outputs = node.run()

    # THEN the node should successfully return the function call value
    assert isinstance(outputs.result, FunctionCall)
    assert outputs.result.name == "test-name"
    assert outputs.result.arguments == {"test-key": "test-value"}
    assert outputs.result.id is None
    assert outputs.log == ""


@pytest.mark.parametrize(
    "code_snippet",
    [
        """
from vellum.client.types.function_call import FunctionCall
def main(input: FunctionCall) -> str:
    return input.type
""",
        """
from vellum.client.types.function_call import FunctionCall
def main(input: FunctionCall) -> str:
    return input["type"]
""",
    ],
)
def test_run_node__function_call_wrapper_type(code_snippet):
    """Test function call wrapper type access using different patterns"""

    # GIVEN a node that accesses the function call type
    class ExampleCodeExecutionNode(CodeExecutionNode[BaseState, str]):
        code = code_snippet
        runtime = "PYTHON_3_11_6"
        code_inputs = {
            "input": FunctionCall(
                name="test-name",
                arguments={
                    "test-key": "test-value",
                },
            )
        }

    # WHEN we run the node
    node = ExampleCodeExecutionNode()
    outputs = node.run()

    # THEN the node should successfully return the function call type
    assert outputs == {"result": "FUNCTION_CALL", "log": ""}


def test_run_node__iter_list():
    # GIVEN a node that will return the first string in a list
    class ExampleCodeExecutionNode(CodeExecutionNode[BaseState, str]):
        code = """\
def main(
    input_list: list
) -> str:
    return next((
        o.value for o in input_list if o.type == "STRING"
    ), None)
"""
        runtime = "PYTHON_3_11_6"
        code_inputs = {
            "input_list": [
                StringVellumValue(value="foo"),
                NumberVellumValue(value=1),
            ]
        }

    # WHEN we run the node
    node = ExampleCodeExecutionNode()
    outputs = node.run()

    # THEN the node should successfully access the string value
    assert outputs == {"result": "foo", "log": ""}


def test_run_node__iter_dict():
    # GIVEN a node that will return the first string in a list
    class ExampleCodeExecutionNode(CodeExecutionNode[BaseState, list[str]]):
        code = """\
def main(input_dict: dict) -> list[str]:
    return [item.value for item in input_dict.values()]
"""
        runtime = "PYTHON_3_11_6"
        code_inputs = {
            "input_dict": {
                "foo": "bar",
                "baz": "qux",
            }
        }

    # WHEN we run the node
    node = ExampleCodeExecutionNode()
    outputs = node.run()

    # THEN the node should successfully access the string value
    assert outputs == {"result": ["bar", "qux"], "log": ""}


def test_run_node__show_clearer_runtime_error_message():
    # GIVEN a node that will return the first string in a list
    class ExampleCodeExecutionNode(CodeExecutionNode[BaseState, list[str]]):
        code = """\
def main(items: list[str]) -> list[str]:
    return first(items)

# Helper function to prove out stack traces
def first(items: list[str]) -> str:
    return items[0]
"""
        runtime = "PYTHON_3_11_6"
        code_inputs = {"items": []}

    # WHEN we run the node
    node = ExampleCodeExecutionNode()
    with pytest.raises(NodeException) as exc_info:
        node.run()

    # THEN the node should successfully access the string value
    assert (
        exc_info.value.message
        == """\
Traceback (most recent call last):
  File "ExampleCodeExecutionNode.code.py", line 2, in main
    return first(items)
  File "ExampleCodeExecutionNode.code.py", line 6, in first
    return items[0]

IndexError: list index out of range
"""
    )


def test_run_node__default_function_call_type():
    # GIVEN a node that will return a FunctionCall
    class ExampleCodeExecutionNode(CodeExecutionNode[BaseState, FunctionCall]):
        code = """\
def main(input: str) -> str:
    return None
"""
        runtime = "PYTHON_3_11_6"
        code_inputs = {"input": "foo"}

    # WHEN we run the node
    node = ExampleCodeExecutionNode()
    outputs = node.run()

    # THEN the node should return default function call
    assert outputs == {"result": FunctionCall(name="", arguments={}), "log": ""}


def test_run_node__array_input_with_primitive_values(vellum_client):
    # GIVEN a node that subclasses CodeExecutionNode that processes an array of primitive values
    class State(BaseState):
        pass

    class ExampleCodeExecutionNode(CodeExecutionNode[State, str]):
        code = """\
from typing import List
def main(numbers: List[str]) -> str:
    return " ".join(numbers)
"""
        runtime = "PYTHON_3_11_6"
        packages = [
            CodeExecutionPackage(
                name="openai",
                version="1.0.0",
            )
        ]

        code_inputs = {
            "numbers": ["6", "4", "2"],
        }

    # AND we know what the Code Execution Node will respond with
    mock_code_execution = CodeExecutorResponse(
        log="",
        output=StringVellumValue(value="6 4 2"),
    )
    vellum_client.execute_code.return_value = mock_code_execution

    # WHEN we run the node
    node = ExampleCodeExecutionNode(state=State())
    outputs = node.run()

    # THEN the node should have produced the outputs we expect
    assert outputs == {"result": "6 4 2", "log": ""}

    # AND we should have invoked the Code with the expected inputs
    vellum_client.execute_code.assert_called_once_with(
        input_values=[
            ArrayInput(
                name="numbers",
                value=[
                    StringVellumValue(value="6"),
                    StringVellumValue(value="4"),
                    StringVellumValue(value="2"),
                ],
            )
        ],
        code="""\
from typing import List
def main(numbers: List[str]) -> str:
    return " ".join(numbers)
""",
        runtime="PYTHON_3_11_6",
        output_type="STRING",
        packages=[
            CodeExecutionPackage(
                name="openai",
                version="1.0.0",
            )
        ],
        request_options=None,
    )


def test_run_node__vellum_secret_forces_api_execution(vellum_client):
    """Test that CodeExecutionNode with VellumSecretReference forces API execution instead of inline execution."""

    # GIVEN a node that subclasses CodeExecutionNode with VellumSecretReference but no packages
    # This should normally run inline, but the presence of secrets should force API execution
    class State(BaseState):
        pass

    class ExampleCodeExecutionNode(CodeExecutionNode[State, str]):
        code = """
def main(secret: str) -> str:
    return f"Secret value: {secret}"
"""
        runtime = "PYTHON_3_11_6"
        # Note: No packages specified, which would normally trigger inline execution

        code_inputs = {
            "secret": VellumSecretReference("MY_SECRET"),
        }

    # AND we know what the Code Execution Node will respond with
    mock_code_execution = CodeExecutorResponse(
        log="",
        output=StringVellumValue(value="Secret value: my_secret_value"),
    )
    vellum_client.execute_code.return_value = mock_code_execution

    # WHEN we run the node
    node = ExampleCodeExecutionNode(state=State())
    outputs = node.run()

    # THEN the node should have produced the outputs we expect
    assert outputs == {"result": "Secret value: my_secret_value", "log": ""}

    # AND we should have invoked the Code via API (not inline) due to the secret
    vellum_client.execute_code.assert_called_once_with(
        input_values=[
            CodeExecutorSecretInput(
                name="secret",
                value="MY_SECRET",
            )
        ],
        code="""
def main(secret: str) -> str:
    return f"Secret value: {secret}"
""",
        runtime="PYTHON_3_11_6",
        output_type="STRING",
        packages=[],
        request_options=None,
    )


def test_run_node__undefined_input_skipped():
    """
    Confirm that when an undefined value is passed as an input, it is skipped rather than raising an error.
    The function should use the default parameter value when undefined input is skipped.
    """

    # GIVEN a node with both a valid input and an undefined input that runs inline
    class State(BaseState):
        pass

    class ExampleCodeExecutionNode(CodeExecutionNode[State, int]):
        code = """\
def main(word: str, undefined_input: str = "default") -> int:
    return len(word) + len(undefined_input)
"""
        runtime = "PYTHON_3_11_6"
        packages = []

        code_inputs = {
            "word": "hello",
            "undefined_input": undefined,
        }

    # WHEN we run the node
    node = ExampleCodeExecutionNode(state=State())
    outputs = node.run()

    # THEN the node should run successfully without raising an error
    assert outputs == {"result": 12, "log": ""}  # len("hello") + len("default") = 5 + 7 = 12


def test_run_node__undefined_input_skipped_api_execution(vellum_client):
    """
    Confirm that when an undefined value is passed as an input in API execution mode,
    it is skipped rather than raising an error.
    """

    # GIVEN a node with both a valid input and an undefined input that forces API execution
    class State(BaseState):
        pass

    class ExampleCodeExecutionNode(CodeExecutionNode[State, int]):
        code = """\
def main(word: str) -> int:
    return len(word)
"""
        runtime = "PYTHON_3_11_6"
        packages = [CodeExecutionPackage(name="requests", version="2.0.0")]

        code_inputs = {
            "word": "hello",
            "undefined_input": undefined,
        }

    # AND we know what the Code Execution Node will respond with
    mock_code_execution = CodeExecutorResponse(
        log="",
        output=NumberVellumValue(value=5),
    )
    vellum_client.execute_code.return_value = mock_code_execution

    # WHEN we run the node
    node = ExampleCodeExecutionNode(state=State())
    outputs = node.run()

    # THEN the node should run successfully without raising an error
    assert outputs == {"result": 5, "log": ""}

    # AND the API should be called with only the non-undefined input
    vellum_client.execute_code.assert_called_once_with(
        input_values=[StringInput(name="word", value="hello")],
        code="""\
def main(word: str) -> int:
    return len(word)
""",
        runtime="PYTHON_3_11_6",
        output_type="NUMBER",
        packages=[CodeExecutionPackage(name="requests", version="2.0.0")],
        request_options=None,
    )
