import pytest
from contextlib import redirect_stdout
from importlib.metadata import version
import io
import json
from unittest import mock
from uuid import uuid4

import requests_mock

from workflow_server.code_exec_runner import run_code_exec_stream
from workflow_server.server import create_app


def flask_stream(request_body: dict) -> tuple[int, list]:
    flask_app = create_app()
    with flask_app.test_client() as test_client:
        response = test_client.post("/workflow/stream", json=request_body)
        status_code = response.status_code

        return status_code, [
            json.loads(line)
            for line in response.data.decode().split("\n")
            if line
            and line
            not in [
                "WAITING",
                "END",
            ]
        ]


def code_exec_stream(request_body: dict) -> tuple[int, list]:
    output = io.StringIO()

    with mock.patch("os.read") as mocked_os_read, redirect_stdout(output):
        mocked_os_read.return_value = (json.dumps(request_body) + "\n--vellum-input-stop--\n").encode("utf-8")
        run_code_exec_stream()

    lines = output.getvalue().split("\n")
    events = []
    for line in lines:
        if "--event--" in line:
            events.append(json.loads(line.replace("--event--", "")))

    return 200, events


@pytest.fixture(params=[flask_stream, code_exec_stream])
def both_stream_types(request):
    return request.param


def test_stream_workflow_flask_route__verify_headers():
    # GIVEN a valid request body
    span_id = uuid4()
    request_body = {
        "execution_id": str(span_id),
        "inputs": [],
        "environment_api_key": "test",
        "module": "workflow",
        "timeout": 360,
        "files": {
            "__init__.py": "",
            "workflow.py": """\
from vellum.workflows import BaseWorkflow

class Workflow(BaseWorkflow):
    pass
""",
        },
    }

    # WHEN we call the stream route
    flask_app = create_app()
    with flask_app.test_client() as test_client:
        response = test_client.post("/workflow/stream", json=request_body)
        status_code = response.status_code

    # THEN we get a 200 response
    assert status_code == 200, response.text

    # AND the version headers are present
    assert "X-Vellum-SDK-Version" in response.headers
    assert "X-Vellum-Server-Version" in response.headers


def test_stream_workflow_route__happy_path(both_stream_types):
    # GIVEN a valid request body
    span_id = uuid4()
    request_body = {
        "execution_id": str(span_id),
        "inputs": [],
        "environment_api_key": "test",
        "module": "workflow",
        "timeout": 360,
        "files": {
            "__init__.py": "",
            "workflow.py": """\
from vellum.workflows import BaseWorkflow

class Workflow(BaseWorkflow):
    class Outputs(BaseWorkflow.Outputs):
        foo = "hello"
""",
        },
    }

    # WHEN we call the stream route
    status_code, events = both_stream_types(request_body)

    # THEN we get a 200 response
    assert status_code == 200, events

    # THEN we get the expected events
    assert events[0] == {
        "id": mock.ANY,
        "trace_id": mock.ANY,
        "span_id": str(span_id),
        "timestamp": mock.ANY,
        "api_version": "2024-10-25",
        "parent": None,
        "name": "vembda.execution.initiated",
        "body": {
            "sdk_version": version("vellum-ai"),
            "server_version": "local",
        },
    }

    assert events[1]["name"] == "workflow.execution.initiated", events[1]
    assert "display_context" in events[1]["body"], events[1]["body"]
    assert events[2]["name"] == "workflow.execution.fulfilled", events[2]

    assert events[3] == {
        "id": mock.ANY,
        "trace_id": events[0]["trace_id"],
        "span_id": str(span_id),
        "timestamp": mock.ANY,
        "api_version": "2024-10-25",
        "parent": None,
        "name": "vembda.execution.fulfilled",
        "body": mock.ANY,
    }
    assert events[3]["body"] == {
        "exit_code": 0,
        "log": "",
        "stderr": "",
        "timed_out": False,
        "container_overhead_latency": mock.ANY,
    }

    assert len(events) == 4


def test_stream_workflow_route__happy_path_with_inputs(both_stream_types):
    # GIVEN a valid request body
    span_id = uuid4()
    request_body = {
        "timeout": 360,
        "execution_id": str(span_id),
        "inputs": [
            {"name": "foo", "type": "STRING", "value": "hello"},
        ],
        "environment_api_key": "test",
        "module": "workflow",
        "files": {
            "__init__.py": "",
            "workflow.py": """\
from vellum.workflows import BaseWorkflow
from vellum.workflows.state import BaseState
from .inputs import Inputs

class Workflow(BaseWorkflow[Inputs, BaseState]):
    class Outputs(BaseWorkflow.Outputs):
        foo = "hello"
""",
            "inputs.py": """\
from vellum.workflows.inputs import BaseInputs

class Inputs(BaseInputs):
    foo: str
""",
        },
    }

    # WHEN we call the stream route
    status_code, events = both_stream_types(request_body)

    # THEN we get a 200 response
    assert status_code == 200, events

    # THEN we get the expected events
    assert events[0] == {
        "id": mock.ANY,
        "trace_id": mock.ANY,
        "span_id": str(span_id),
        "timestamp": mock.ANY,
        "api_version": "2024-10-25",
        "parent": None,
        "name": "vembda.execution.initiated",
        "body": {
            "sdk_version": version("vellum-ai"),
            "server_version": "local",
        },
    }

    assert events[1]["name"] == "workflow.execution.initiated", events[1]
    assert "display_context" in events[1]["body"], events[1]["body"]
    assert events[2]["name"] == "workflow.execution.fulfilled", events[2]

    assert events[3] == {
        "id": mock.ANY,
        "trace_id": events[0]["trace_id"],
        "span_id": str(span_id),
        "timestamp": mock.ANY,
        "api_version": "2024-10-25",
        "parent": None,
        "name": "vembda.execution.fulfilled",
        "body": mock.ANY,
    }
    assert events[3]["body"] == {
        "exit_code": 0,
        "log": "",
        "stderr": "",
        "timed_out": False,
        "container_overhead_latency": mock.ANY,
    }

    assert len(events) == 4


def test_stream_workflow_route__happy_path_with_state(both_stream_types):
    # GIVEN a valid request body
    span_id = uuid4()
    request_body = {
        "timeout": 360,
        "execution_id": str(span_id),
        "state": {"foo": "bar"},
        "environment_api_key": "test",
        "module": "workflow",
        "files": {
            "__init__.py": "",
            "workflow.py": """\
from vellum.workflows import BaseWorkflow
from vellum.workflows.inputs import BaseInputs
from .state import State

class Workflow(BaseWorkflow[BaseInputs, State]):
    class Outputs(BaseWorkflow.Outputs):
        foo = State.foo
""",
            "state.py": """\
from vellum.workflows.state import BaseState

class State(BaseState):
    foo: str
""",
        },
    }

    # WHEN we call the stream route
    status_code, events = both_stream_types(request_body)

    # THEN we get a 200 response
    assert status_code == 200, events

    # THEN we get the expected events
    assert events[0] == {
        "id": mock.ANY,
        "trace_id": mock.ANY,
        "span_id": str(span_id),
        "timestamp": mock.ANY,
        "api_version": "2024-10-25",
        "parent": None,
        "name": "vembda.execution.initiated",
        "body": {
            "sdk_version": version("vellum-ai"),
            "server_version": "local",
        },
    }

    assert events[1]["name"] == "workflow.execution.initiated", events[1]
    assert "display_context" in events[1]["body"], events[1]["body"]
    assert events[2]["name"] == "workflow.execution.fulfilled", events[2]
    assert events[2]["body"]["outputs"] == {"foo": "bar"}

    assert events[3] == {
        "id": mock.ANY,
        "trace_id": events[0]["trace_id"],
        "span_id": str(span_id),
        "timestamp": mock.ANY,
        "api_version": "2024-10-25",
        "parent": None,
        "name": "vembda.execution.fulfilled",
        "body": mock.ANY,
    }
    assert events[3]["body"] == {
        "exit_code": 0,
        "log": "",
        "stderr": "",
        "timed_out": False,
        "container_overhead_latency": mock.ANY,
    }

    assert len(events) == 4


def test_stream_workflow_route__bad_indent_in_inputs_file(both_stream_types):
    # GIVEN a valid request body
    span_id = uuid4()
    request_body = {
        "timeout": 360,
        "execution_id": str(span_id),
        "inputs": [
            {"name": "foo", "type": "STRING", "value": "hello"},
        ],
        "environment_api_key": "test",
        "module": "workflow",
        "files": {
            "__init__.py": "",
            "workflow.py": """\
from vellum.workflows import BaseWorkflow
from vellum.workflows.state import BaseState
from .inputs import Inputs

class Workflow(BaseWorkflow[Inputs, BaseState]):
    class Outputs(BaseWorkflow.Outputs):
        foo = "hello"
""",
            "inputs.py": """\
from vellum.workflows.inputs import BaseInputs

  class Inputs(BaseInputs):
     foo: str
""",
        },
    }

    # WHEN we call the stream route
    status_code, events = both_stream_types(request_body)

    # THEN we get a 200 response
    assert status_code == 200, events

    # THEN we get the expected events
    assert events[0] == {
        "id": mock.ANY,
        "trace_id": mock.ANY,
        "span_id": str(span_id),
        "timestamp": mock.ANY,
        "api_version": "2024-10-25",
        "parent": None,
        "name": "vembda.execution.initiated",
        "body": {
            "sdk_version": version("vellum-ai"),
            "server_version": "local",
        },
    }

    assert events[1] == {
        "id": mock.ANY,
        "trace_id": events[0]["trace_id"],
        "span_id": str(span_id),
        "timestamp": mock.ANY,
        "api_version": "2024-10-25",
        "parent": None,
        "name": "vembda.execution.fulfilled",
        "body": mock.ANY,
    }
    assert events[1]["body"] == {
        "exit_code": -1,
        "log": "",
        "stderr": "Failed to initialize workflow: unexpected indent (inputs.py, line 3)",
        "timed_out": False,
        "container_overhead_latency": mock.ANY,
    }

    assert len(events) == 2


def test_stream_workflow_route__cancel(both_stream_types):
    # GIVEN a valid request body
    span_id = uuid4()
    request_body = {
        "timeout": 360,
        "execution_id": str(span_id),
        "inputs": [],
        "environment_api_key": "test",
        "module": "workflow",
        "vembda_public_url": "http://test.biz",
        "files": {
            "__init__.py": "",
            "workflow.py": """\
import time

from vellum.workflows.nodes.bases.base import BaseNode
from vellum.workflows.workflows.base import BaseWorkflow


class StartNode(BaseNode):
    class Outputs(BaseNode.Outputs):
        value: str

    def run(self) -> Outputs:
        time.sleep(2)
        return self.Outputs(value="hello world")


class BasicCancellableWorkflow(BaseWorkflow):
    graph = StartNode
    class Outputs(BaseWorkflow.Outputs):
        final_value = StartNode.Outputs.value

""",
        },
    }

    # WHEN we call the stream route with a mock cancelled return true
    with requests_mock.Mocker() as mocker:
        mocker.get(
            f"http://test.biz/vembda-public/cancel-workflow-execution-status/{span_id}", json={"cancelled": True}
        )
        status_code, events = both_stream_types(request_body)

    # THEN we get a 200 response
    assert status_code == 200, events

    # THEN we get the expected cancelled events
    assert events[0] == {
        "id": mock.ANY,
        "trace_id": mock.ANY,
        "span_id": str(span_id),
        "timestamp": mock.ANY,
        "api_version": "2024-10-25",
        "parent": None,
        "name": "vembda.execution.initiated",
        "body": {
            "sdk_version": version("vellum-ai"),
            "server_version": "local",
        },
    }

    assert events[1]["name"] == "workflow.execution.initiated", events[1]
    assert "display_context" in events[1]["body"], events[1]["body"]
    cancelled_event = None
    for event in events:
        if event["name"] == "workflow.execution.rejected":
            cancelled_event = event
            break

    assert cancelled_event["body"]["error"]["message"] == "Workflow run cancelled"


def test_stream_workflow_route__very_large_events(both_stream_types):
    # GIVEN a valid request body
    span_id = uuid4()
    request_body = {
        "execution_id": str(span_id),
        "inputs": [
            {"name": "foo", "type": "STRING", "value": "hello" * 10_000_000},
        ],
        "environment_api_key": "test",
        "module": "workflow",
        "files": {
            "__init__.py": "",
            "workflow.py": """\
from vellum.workflows import BaseWorkflow
from vellum.workflows.state import BaseState
from .inputs import Inputs

class Workflow(BaseWorkflow[Inputs, BaseState]):
    class Outputs(BaseWorkflow.Outputs):
        foo = "hello"
""",
            "inputs.py": """\
from vellum.workflows.inputs import BaseInputs

class Inputs(BaseInputs):
    foo: str
""",
        },
    }

    # WHEN we call the stream route
    status_code, events = both_stream_types(request_body)

    # THEN we get a 200 response
    assert status_code == 200, events

    # THEN we get the expected events
    assert events[0] == {
        "id": mock.ANY,
        "trace_id": mock.ANY,
        "span_id": str(span_id),
        "timestamp": mock.ANY,
        "api_version": "2024-10-25",
        "parent": None,
        "name": "vembda.execution.initiated",
        "body": {
            "sdk_version": version("vellum-ai"),
            "server_version": "local",
        },
    }

    assert events[1]["name"] == "workflow.execution.initiated", events[1]
    assert "display_context" in events[1]["body"], events[1]["body"]
    assert events[2]["name"] == "workflow.execution.fulfilled", events[2]

    assert events[3] == {
        "id": mock.ANY,
        "trace_id": events[0]["trace_id"],
        "span_id": str(span_id),
        "timestamp": mock.ANY,
        "api_version": "2024-10-25",
        "parent": None,
        "name": "vembda.execution.fulfilled",
        "body": mock.ANY,
    }
    assert events[3]["body"] == {
        "exit_code": 0,
        "log": "",
        "stderr": "",
        "timed_out": False,
        "container_overhead_latency": mock.ANY,
    }

    assert len(events) == 4


def test_stream_workflow_route__happy_path_run_from_node(both_stream_types):
    # GIVEN a valid request body representing a run from a node
    node_id = uuid4()
    span_id = uuid4()
    request_body = {
        "timeout": 360,
        "execution_id": str(span_id),
        "node_id": str(node_id),
        "environment_api_key": "test",
        "module": "workflow",
        "files": {
            "__init__.py": "from .display import *",
            "workflow.py": """\
from vellum.workflows import BaseWorkflow
from .nodes.start import StartNode
from .nodes.end import EndNode

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

    class Outputs(BaseWorkflow.Outputs):
        foo = EndNode.Outputs.value
""",
            "nodes/__init__.py": """
from .start import StartNode
from .end import EndNode

__all__ = ["StartNode", "EndNode"]
""",
            "nodes/start.py": """\
from vellum.workflows.nodes import BaseNode

class StartNode(BaseNode):
    class Outputs(BaseNode.Outputs):
        value = "apple"
""",
            "nodes/end.py": """\
from vellum.workflows.nodes import BaseNode
from .start import StartNode

class EndNode(BaseNode):
    class Outputs(BaseNode.Outputs):
        value = StartNode.Outputs.value.coalesce("banana")
""",
            "display/__init__.py": """
from .nodes import *
from .workflow import *
""",
            "display/workflow.py": """
from uuid import UUID
from vellum_ee.workflows.display.workflows import BaseWorkflowDisplay
from vellum_ee.workflows.display.base import EdgeDisplay
from ..nodes.start import StartNode
from ..nodes.end import EndNode

class WorkflowDisplay(BaseWorkflowDisplay):
    edge_displays = {
        (StartNode.Ports.default, EndNode): EdgeDisplay(
            id=UUID("63606ff1-1c70-4516-92c6-cbaba9336424")
        ),
    }
""",
            "display/nodes/__init__.py": """
from .start import StartNodeDisplay
from .end import EndNodeDisplay

__all__ = ["StartNodeDisplay", "EndNodeDisplay"]
""",
            "display/nodes/start.py": """\
from vellum_ee.workflows.display.nodes import BaseNodeDisplay
from ...nodes.start import StartNode

class StartNodeDisplay(BaseNodeDisplay[StartNode]):
    pass
""",
            "display/nodes/end.py": f"""\
from uuid import UUID

from vellum_ee.workflows.display.nodes import BaseNodeDisplay
from ...nodes.end import EndNode

class EndNodeDisplay(BaseNodeDisplay[EndNode]):
    node_id = UUID("{node_id}")
""",
        },
    }

    # WHEN we call the stream route
    status_code, events = both_stream_types(request_body)

    # THEN we get a 200 response
    assert status_code == 200, events

    # THEN we get the expected workflow and node fulfilled events
    assert events[-4]["name"] == "node.execution.fulfilled", json.dumps(events[-1])
    assert events[-4]["body"]["node_definition"]["id"] == str(node_id)
    assert events[-2]["name"] == "workflow.execution.fulfilled", json.dumps(events[-1])
    assert events[-2]["body"]["outputs"] == {"foo": "banana"}
    assert [event["name"] for event in events] == [
        "vembda.execution.initiated",
        "workflow.execution.initiated",
        "node.execution.initiated",
        "workflow.execution.snapshotted",
        "node.execution.fulfilled",
        "workflow.execution.streaming",
        "workflow.execution.fulfilled",
        "vembda.execution.fulfilled",
    ]


def test_stream_workflow_route__happy_path_run_from_node_with_state(both_stream_types):
    # GIVEN a valid request body representing a run from a node with state
    node_id = uuid4()
    span_id = uuid4()
    request_body = {
        "timeout": 360,
        "execution_id": str(span_id),
        "node_id": str(node_id),
        "environment_api_key": "test",
        "module": "workflow",
        "state": {"meta": {"node_outputs": {"StartNode.Outputs.value": "cherry"}}},
        "files": {
            "__init__.py": "from .display import *",
            "workflow.py": """\
from vellum.workflows import BaseWorkflow
from .nodes.start import StartNode
from .nodes.end import EndNode

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

    class Outputs(BaseWorkflow.Outputs):
        foo = EndNode.Outputs.value
""",
            "nodes/__init__.py": """
from .start import StartNode
from .end import EndNode

__all__ = ["StartNode", "EndNode"]
""",
            "nodes/start.py": """\
import random
from vellum.workflows.nodes import BaseNode

class StartNode(BaseNode):
    class Outputs(BaseNode.Outputs):
        value = str
    def run(self) -> Outputs:
        return self.Outputs(value=random.choice(["apple", "banana", "cherry"]))
""",
            "nodes/end.py": """\
from vellum.workflows.nodes import BaseNode
from .start import StartNode

class EndNode(BaseNode):
    fruit = StartNode.Outputs.value.coalesce("date")

    class Outputs(BaseNode.Outputs):
        value: str

    def run(self) -> Outputs:
        return self.Outputs(value=self.fruit)
""",
            "display/__init__.py": """
from .nodes import *
from .workflow import *
""",
            "display/workflow.py": """
from uuid import UUID
from vellum_ee.workflows.display.workflows import BaseWorkflowDisplay
from vellum_ee.workflows.display.base import EdgeDisplay
from ..nodes.start import StartNode
from ..nodes.end import EndNode

class WorkflowDisplay(BaseWorkflowDisplay):
    edge_displays = {
        (StartNode.Ports.default, EndNode): EdgeDisplay(
            id=UUID("63606ff1-1c70-4516-92c6-cbaba9336424")
        ),
    }
""",
            "display/nodes/__init__.py": """
from .start import StartNodeDisplay
from .end import EndNodeDisplay

__all__ = ["StartNodeDisplay", "EndNodeDisplay"]
""",
            "display/nodes/start.py": """\
from vellum_ee.workflows.display.nodes import BaseNodeDisplay
from ...nodes.start import StartNode

class StartNodeDisplay(BaseNodeDisplay[StartNode]):
    pass
""",
            "display/nodes/end.py": f"""\
from uuid import UUID

from vellum_ee.workflows.display.nodes import BaseNodeDisplay
from ...nodes.end import EndNode

class EndNodeDisplay(BaseNodeDisplay[EndNode]):
    node_id = UUID("{node_id}")
""",
        },
    }

    # WHEN we call the stream route
    status_code, events = both_stream_types(request_body)

    # THEN we get a 200 response
    assert status_code == 200, events

    # THEN we get the expected workflow and node initiated events
    assert len(events) > 3, json.dumps(events[-1])
    assert events[2]["name"] == "node.execution.initiated", json.dumps(events[-1])
    assert events[2]["body"]["inputs"] == {"fruit": "cherry"}


def test_stream_workflow_route__with_environment_variables(both_stream_types):
    # GIVEN a valid request body with environment variables
    span_id = uuid4()
    request_body = {
        "execution_id": str(span_id),
        "inputs": [],
        "environment_api_key": "test",
        "module": "workflow",
        "timeout": 360,
        "environment_variables": {"TEST_ENV_VAR": "test_value", "ANOTHER_VAR": "another_value"},
        "files": {
            "__init__.py": "",
            "workflow.py": """\
from vellum.workflows import BaseWorkflow
from vellum.workflows.references import EnvironmentVariableReference

class Workflow(BaseWorkflow):
    class Outputs(BaseWorkflow.Outputs):
        env_var_value = EnvironmentVariableReference(name="TEST_ENV_VAR", default="not_found")
        another_var_value = EnvironmentVariableReference(name="ANOTHER_VAR", default="not_found")
""",
        },
    }

    # WHEN we call the stream route
    status_code, events = both_stream_types(request_body)

    # THEN we get a 200 response
    assert status_code == 200, events

    # THEN we get the expected events
    assert events[0]["name"] == "vembda.execution.initiated"
    assert events[1]["name"] == "workflow.execution.initiated"
    assert events[2]["name"] == "workflow.execution.fulfilled"

    # AND the environment variables are accessible in the workflow
    outputs = events[2]["body"]["outputs"]
    assert outputs["env_var_value"] == "test_value"
    assert outputs["another_var_value"] == "another_value"
