# tests/test_hooks.py
"""Tests for context lifecycle hooks."""

from collections.abc import Mapping
from functools import cached_property
from typing import ClassVar, Literal
from unittest.mock import AsyncMock

import pytest

from workflow_engine import (
    Context,
    Data,
    DataMapping,
    Edge,
    ExecutionAlgorithm,
    Node,
    NodeTypeInfo,
    Params,
    ShouldYield,
    StringValue,
    Workflow,
    WorkflowErrors,
    WorkflowExecutionResult,
    WorkflowExecutionResultStatus,
)
from workflow_engine.contexts import InMemoryContext
from workflow_engine.core.io import InputNode, OutputNode
from workflow_engine.execution import TopologicalExecutionAlgorithm
from workflow_engine.execution.parallel import (
    ErrorHandlingMode,
    ParallelExecutionAlgorithm,
)
from workflow_engine.nodes import ConstantStringNode, ErrorNode


@pytest.fixture(params=["topological", "parallel"])
def algorithm(request) -> ExecutionAlgorithm:
    if request.param == "topological":
        return TopologicalExecutionAlgorithm()
    elif request.param == "parallel":
        return ParallelExecutionAlgorithm()
    else:
        raise ValueError(f"Invalid algorithm: {request.param}")


def _simple_workflow() -> Workflow:
    """A minimal workflow: (no input) -> constant -> output."""
    input_node = InputNode.empty()
    output_node = OutputNode.from_fields(value=StringValue)
    constant = ConstantStringNode.from_value(id="constant", value="hello")
    return Workflow(
        input_node=input_node,
        output_node=output_node,
        inner_nodes=[constant],
        edges=[
            Edge.from_nodes(
                source=constant,
                source_key="value",
                target=output_node,
                target_key="value",
            )
        ],
    )


def _error_workflow() -> Workflow:
    """A workflow that raises an error during execution."""
    input_node = InputNode.empty()
    output_node = OutputNode.from_fields(value=StringValue)
    constant = ConstantStringNode.from_value(id="constant", value="hello")
    error = ErrorNode.from_name(id="error", name="TestError")
    return Workflow(
        input_node=input_node,
        output_node=output_node,
        inner_nodes=[constant, error],
        edges=[
            Edge.from_nodes(
                source=constant,
                source_key="value",
                target=error,
                target_key="info",
            ),
            Edge.from_nodes(
                source=constant,
                source_key="value",
                target=output_node,
                target_key="value",
            ),
        ],
    )


class ExpandingOutput(Data):
    value: StringValue


# A node that expands into a subgraph (returns a Workflow from run())
# Its output_type matches the inner workflow's output schema.
class ExpandingNode(Node[Data, ExpandingOutput, Params]):
    TYPE_INFO: ClassVar[NodeTypeInfo] = NodeTypeInfo.from_parameter_type(
        name="Expanding",
        display_name="Expanding",
        description="A node that expands into a subgraph.",
        version="1.0.0",
        parameter_type=Params,
    )
    type: Literal["Expanding"] = "Expanding"  # pyright: ignore[reportIncompatibleVariableOverride]

    output_value: str

    @cached_property
    def input_type(self):
        return Data

    @cached_property
    def output_type(self):
        return ExpandingOutput

    async def run(self, context: Context, input: Data) -> Workflow:
        inner_input = InputNode.empty()
        inner_output = OutputNode.from_fields(value=StringValue)
        constant = ConstantStringNode.from_value(
            id="inner_constant", value=self.output_value
        )
        return Workflow(
            input_node=inner_input,
            output_node=inner_output,
            inner_nodes=[constant],
            edges=[
                Edge.from_nodes(
                    source=constant,
                    source_key="value",
                    target=inner_output,
                    target_key="value",
                )
            ],
        )

    @classmethod
    def create(cls, id: str, output_value: str = "expanded") -> "ExpandingNode":
        return cls(id=id, params=Params(), output_value=output_value)


class YieldingNode(Node[Data, ExpandingOutput, Params]):
    TYPE_INFO: ClassVar[NodeTypeInfo] = NodeTypeInfo.from_parameter_type(
        name="HookYielding",
        display_name="HookYielding",
        description="Always raises ShouldYield.",
        version="1.0.0",
        parameter_type=Params,
    )
    type: Literal["HookYielding"] = "HookYielding"  # pyright: ignore[reportIncompatibleVariableOverride]

    @cached_property
    def input_type(self):
        return Data

    @cached_property
    def output_type(self):
        return ExpandingOutput

    async def run(self, context: Context, input: Data) -> ExpandingOutput:
        raise ShouldYield("waiting for approval")


def _error_and_yield_workflow() -> Workflow:
    """A workflow where one independent branch errors and another yields."""
    input_node = InputNode.empty()
    output_node = OutputNode.from_fields(value=StringValue)
    constant = ConstantStringNode.from_value(id="constant", value="hello")
    error = ErrorNode.from_name(id="error", name="TestError")
    yielding = YieldingNode(id="yielding", params=Params())
    return Workflow(
        input_node=input_node,
        output_node=output_node,
        inner_nodes=[constant, error, yielding],
        edges=[
            Edge.from_nodes(
                source=constant,
                source_key="value",
                target=error,
                target_key="info",
            ),
            Edge.from_nodes(
                source=yielding,
                source_key="value",
                target=output_node,
                target_key="value",
            ),
        ],
    )


def _yielding_workflow() -> Workflow:
    yielding = YieldingNode(id="yielding", params=Params())
    input_node = InputNode.empty()
    output_node = OutputNode.from_fields(value=StringValue)
    return Workflow(
        input_node=input_node,
        output_node=output_node,
        inner_nodes=[yielding],
        edges=[
            Edge.from_nodes(
                source=yielding,
                source_key="value",
                target=output_node,
                target_key="value",
            )
        ],
    )


def _expanding_workflow(output_value: str = "expanded") -> Workflow:
    expanding = ExpandingNode.create(id="expanding", output_value=output_value)
    input_node = InputNode.empty()
    output_node = OutputNode.from_fields(value=StringValue)
    return Workflow(
        input_node=input_node,
        output_node=output_node,
        inner_nodes=[expanding],
        edges=[
            Edge.from_nodes(
                source=expanding,
                source_key="value",
                target=output_node,
                target_key="value",
            )
        ],
    )


class TestOnNodeStart:
    @pytest.mark.asyncio
    async def test_called_for_each_node(self):
        workflow = _simple_workflow()
        context = InMemoryContext()
        mock = AsyncMock(side_effect=context.on_node_start)
        context.on_node_start = mock

        algorithm = TopologicalExecutionAlgorithm()
        result = await algorithm.execute(context=context, workflow=workflow, input={})

        assert result.status is WorkflowExecutionResultStatus.SUCCESS
        assert mock.call_count == 3  # input_node, constant, output_node
        called_ids = {call.kwargs["node"].id for call in mock.call_args_list}
        assert called_ids == {
            workflow.input_node.id,
            "constant",
            workflow.output_node.id,
        }

    @pytest.mark.asyncio
    async def test_can_skip_node_execution(self):
        """Returning a DataMapping from on_node_start bypasses run()."""
        workflow = _simple_workflow()
        context = InMemoryContext()
        original = context.on_node_start

        async def on_node_start(
            *, node: Node, input: DataMapping
        ) -> DataMapping | None:
            if node.id == "constant":
                return {"value": StringValue("cached")}
            return await original(node=node, input=input)

        context.on_node_start = on_node_start

        algorithm = TopologicalExecutionAlgorithm()
        result = await algorithm.execute(context=context, workflow=workflow, input={})

        assert result.status is WorkflowExecutionResultStatus.SUCCESS
        assert result.output == {"value": StringValue("cached")}


class TestOnNodeFinish:
    @pytest.mark.asyncio
    async def test_called_for_each_node(self):
        workflow = _simple_workflow()
        context = InMemoryContext()
        mock = AsyncMock(side_effect=context.on_node_finish)
        context.on_node_finish = mock

        algorithm = TopologicalExecutionAlgorithm()
        result = await algorithm.execute(context=context, workflow=workflow, input={})

        assert result.status is WorkflowExecutionResultStatus.SUCCESS
        assert mock.call_count == 3  # input_node, constant, output_node
        called_ids = {call.kwargs["node"].id for call in mock.call_args_list}
        assert called_ids == {
            workflow.input_node.id,
            "constant",
            workflow.output_node.id,
        }

    @pytest.mark.asyncio
    async def test_can_modify_output(self):
        """The context can transform a node's output via on_node_finish."""
        workflow = _simple_workflow()
        context = InMemoryContext()

        async def on_node_finish(*, node, input, output):
            if node.id == "constant":
                return {"value": StringValue("overridden")}
            return output

        context.on_node_finish = on_node_finish

        algorithm = TopologicalExecutionAlgorithm()
        result = await algorithm.execute(context=context, workflow=workflow, input={})

        assert result.status is WorkflowExecutionResultStatus.SUCCESS
        assert result.output == {"value": StringValue("overridden")}

    @pytest.mark.asyncio
    async def test_not_called_for_expanding_nodes(self):
        """on_node_finish is not called when a node returns a Workflow."""
        workflow = _expanding_workflow()
        context = InMemoryContext()
        finish_ids: list[str] = []

        async def on_node_finish(*, node, input, output):
            finish_ids.append(node.id)
            return output

        context.on_node_finish = on_node_finish

        algorithm = TopologicalExecutionAlgorithm()
        result = await algorithm.execute(context=context, workflow=workflow, input={})

        assert result.status is WorkflowExecutionResultStatus.SUCCESS
        assert "expanding" not in finish_ids


class TestOnNodeExpand:
    @pytest.mark.asyncio
    async def test_called_when_node_emits_workflow(self):
        workflow = _expanding_workflow()
        context = InMemoryContext()
        mock = AsyncMock(side_effect=context.on_node_expand)
        context.on_node_expand = mock

        algorithm = TopologicalExecutionAlgorithm()
        result = await algorithm.execute(context=context, workflow=workflow, input={})

        assert result.status is WorkflowExecutionResultStatus.SUCCESS
        assert result.output == {"value": StringValue("expanded")}
        mock.assert_called_once()
        call_kwargs = mock.call_args.kwargs
        assert call_kwargs["node"].id == "expanding"
        assert isinstance(call_kwargs["workflow"], Workflow)

    @pytest.mark.asyncio
    async def test_can_modify_emitted_workflow(self):
        """The context can replace the emitted workflow via on_node_expand."""
        workflow = _expanding_workflow()
        context = InMemoryContext()

        async def on_node_expand(*, node, input, workflow):
            inner_input = InputNode.empty()
            inner_output = OutputNode.from_fields(value=StringValue)
            constant = ConstantStringNode.from_value(
                id="replaced_constant", value="replaced"
            )
            return Workflow(
                input_node=inner_input,
                output_node=inner_output,
                inner_nodes=[constant],
                edges=[
                    Edge.from_nodes(
                        source=constant,
                        source_key="value",
                        target=inner_output,
                        target_key="value",
                    )
                ],
            )

        context.on_node_expand = on_node_expand

        algorithm = TopologicalExecutionAlgorithm()
        result = await algorithm.execute(context=context, workflow=workflow, input={})

        assert result.status is WorkflowExecutionResultStatus.SUCCESS
        assert result.output == {"value": StringValue("replaced")}


class TestOnNodeYield:
    @pytest.mark.asyncio
    async def test_called_when_node_yields(
        self,
        algorithm: ExecutionAlgorithm,
    ):
        workflow = _yielding_workflow()
        context = InMemoryContext()
        mock = AsyncMock()
        context.on_node_yield = mock

        result = await algorithm.execute(context=context, workflow=workflow, input={})
        assert result.status is WorkflowExecutionResultStatus.YIELDED

        mock.assert_called_once()

    @pytest.mark.asyncio
    async def test_called_with_correct_arguments(
        self,
        algorithm: ExecutionAlgorithm,
    ):
        workflow = _yielding_workflow()
        context = InMemoryContext()
        mock = AsyncMock()
        context.on_node_yield = mock

        result = await algorithm.execute(context=context, workflow=workflow, input={})
        assert result.status is WorkflowExecutionResultStatus.YIELDED

        kwargs = mock.call_args.kwargs
        assert kwargs["node"].id == "yielding"
        assert isinstance(kwargs["exception"], ShouldYield)
        assert kwargs["exception"].message == "waiting for approval"

    @pytest.mark.asyncio
    async def test_not_called_for_regular_error(
        self,
        algorithm: ExecutionAlgorithm,
    ):
        """on_node_yield must not fire when a node raises a plain error."""
        workflow = _error_workflow()
        context = InMemoryContext()
        mock = AsyncMock()
        context.on_node_yield = mock

        await algorithm.execute(context=context, workflow=workflow, input={})

        mock.assert_not_called()

    @pytest.mark.asyncio
    async def test_not_called_on_success(
        self,
        algorithm: ExecutionAlgorithm,
    ):
        workflow = _simple_workflow()
        context = InMemoryContext()
        mock = AsyncMock()
        context.on_node_yield = mock

        result = await algorithm.execute(context=context, workflow=workflow, input={})
        assert result.status is WorkflowExecutionResultStatus.SUCCESS
        mock.assert_not_called()


class TestOnWorkflowStart:
    @pytest.mark.asyncio
    async def test_called_at_start(self):
        workflow = _simple_workflow()
        context = InMemoryContext()
        mock = AsyncMock(side_effect=context.on_workflow_start)
        context.on_workflow_start = mock

        algorithm = TopologicalExecutionAlgorithm()
        result = await algorithm.execute(context=context, workflow=workflow, input={})

        assert result.status is WorkflowExecutionResultStatus.SUCCESS
        mock.assert_called_once()
        assert mock.call_args.kwargs["workflow"] is workflow

    @pytest.mark.asyncio
    async def test_can_skip_execution(self):
        """Returning a result from on_workflow_start bypasses all node execution."""
        workflow = _simple_workflow()
        context = InMemoryContext()
        cached_output = {"value": StringValue("cached_workflow")}

        async def on_workflow_start(*, workflow, input):
            return WorkflowExecutionResult.success(cached_output)

        context.on_workflow_start = on_workflow_start

        node_start_ids: list[str] = []
        original_on_node_start = context.on_node_start

        async def tracking_on_node_start(*, node, input):
            node_start_ids.append(node.id)
            return await original_on_node_start(node=node, input=input)

        context.on_node_start = tracking_on_node_start

        algorithm = TopologicalExecutionAlgorithm()
        result = await algorithm.execute(context=context, workflow=workflow, input={})

        assert result.status is WorkflowExecutionResultStatus.SUCCESS
        assert result.output == cached_output
        assert node_start_ids == []  # no nodes ran


class TestOnWorkflowFinish:
    @pytest.mark.asyncio
    async def test_called_on_success(self):
        workflow = _simple_workflow()
        context = InMemoryContext()
        mock = AsyncMock(side_effect=context.on_workflow_finish)
        context.on_workflow_finish = mock

        algorithm = TopologicalExecutionAlgorithm()
        result = await algorithm.execute(context=context, workflow=workflow, input={})

        assert result.status is WorkflowExecutionResultStatus.SUCCESS
        mock.assert_called_once()
        assert mock.call_args.kwargs["workflow"] is workflow
        assert mock.call_args.kwargs["output"] == {"value": StringValue("hello")}

    @pytest.mark.asyncio
    async def test_can_modify_output(self):
        workflow = _simple_workflow()
        context = InMemoryContext()

        async def on_workflow_finish(*, workflow, input, output):
            return WorkflowExecutionResult.success({"value": StringValue("modified")})

        context.on_workflow_finish = on_workflow_finish

        algorithm = TopologicalExecutionAlgorithm()
        result = await algorithm.execute(context=context, workflow=workflow, input={})

        assert result.status is WorkflowExecutionResultStatus.SUCCESS
        assert result.output == {"value": StringValue("modified")}

    @pytest.mark.asyncio
    async def test_not_called_on_error(self):
        workflow = _error_workflow()
        context = InMemoryContext()
        mock = AsyncMock()
        context.on_workflow_finish = mock

        algorithm = TopologicalExecutionAlgorithm()
        result = await algorithm.execute(context=context, workflow=workflow, input={})

        assert result.status is WorkflowExecutionResultStatus.ERROR
        mock.assert_not_called()


class TestOnWorkflowYield:
    @pytest.mark.asyncio
    async def test_called_when_workflow_yields(
        self,
        algorithm: ExecutionAlgorithm,
    ):
        workflow = _yielding_workflow()
        context = InMemoryContext()
        mock = AsyncMock(
            return_value=WorkflowExecutionResult.yielded(
                partial_output={},
                node_yields={"yielding": "waiting for approval"},
            )
        )
        context.on_workflow_yield = mock

        result = await algorithm.execute(context=context, workflow=workflow, input={})
        assert result.status is WorkflowExecutionResultStatus.YIELDED

        mock.assert_called_once()

    @pytest.mark.asyncio
    async def test_called_with_correct_arguments(
        self,
        algorithm: ExecutionAlgorithm,
    ):
        workflow = _yielding_workflow()
        context = InMemoryContext()
        mock = AsyncMock(
            return_value=WorkflowExecutionResult.yielded(
                partial_output={},
                node_yields={"yielding": "waiting for approval"},
            )
        )
        context.on_workflow_yield = mock

        result = await algorithm.execute(context=context, workflow=workflow, input={})
        assert result.status is WorkflowExecutionResultStatus.YIELDED

        kwargs = mock.call_args.kwargs
        assert kwargs["workflow"] is workflow
        assert kwargs["partial_output"] == {}
        assert kwargs["node_yields"] == {"yielding": "waiting for approval"}

    @pytest.mark.asyncio
    async def test_not_called_on_success(
        self,
        algorithm: ExecutionAlgorithm,
    ):
        workflow = _simple_workflow()
        context = InMemoryContext()
        mock = AsyncMock()
        context.on_workflow_yield = mock

        result = await algorithm.execute(
            context=context,
            workflow=workflow,
            input={},
        )

        assert result.status is WorkflowExecutionResultStatus.SUCCESS
        mock.assert_not_called()

    @pytest.mark.asyncio
    async def test_not_called_on_error(
        self,
        algorithm: ExecutionAlgorithm,
    ):
        workflow = _error_workflow()
        context = InMemoryContext()
        mock = AsyncMock()
        context.on_workflow_yield = mock

        await algorithm.execute(context=context, workflow=workflow, input={})

        mock.assert_not_called()


class TestOnWorkflowError:
    @pytest.mark.asyncio
    async def test_called_on_error(self):
        workflow = _error_workflow()
        context = InMemoryContext()
        mock = AsyncMock(side_effect=context.on_workflow_error)
        context.on_workflow_error = mock

        algorithm = TopologicalExecutionAlgorithm()
        result = await algorithm.execute(context=context, workflow=workflow, input={})

        assert result.status is WorkflowExecutionResultStatus.ERROR
        mock.assert_called_once()
        assert mock.call_args.kwargs["workflow"] is workflow

    @pytest.mark.asyncio
    async def test_can_modify_errors(self):
        """The context can clear errors via on_workflow_error."""
        workflow = _error_workflow()
        context = InMemoryContext()

        async def on_workflow_error(
            *,
            workflow: Workflow,
            input: DataMapping,
            errors: WorkflowErrors,
            partial_output: DataMapping,
            node_yields: Mapping[str, str],
        ):
            return WorkflowExecutionResult.success(partial_output)

        context.on_workflow_error = on_workflow_error

        algorithm = TopologicalExecutionAlgorithm()
        result = await algorithm.execute(context=context, workflow=workflow, input={})

        assert result.status is WorkflowExecutionResultStatus.SUCCESS

    @pytest.mark.asyncio
    async def test_not_called_on_success(self):
        workflow = _simple_workflow()
        context = InMemoryContext()
        mock = AsyncMock()
        context.on_workflow_error = mock

        algorithm = TopologicalExecutionAlgorithm()
        result = await algorithm.execute(context=context, workflow=workflow, input={})

        assert result.status is WorkflowExecutionResultStatus.SUCCESS
        mock.assert_not_called()


class TestWorkflowYieldPartialOutput:
    @pytest.mark.asyncio
    async def test_partial_output_passed_to_on_workflow_yield(
        self, algorithm: ExecutionAlgorithm
    ):
        """on_workflow_yield receives the partial output from completed nodes."""
        workflow = _yielding_workflow()
        context = InMemoryContext()
        received: dict = {}

        async def on_workflow_yield(
            *,
            workflow: Workflow,
            input: DataMapping,
            partial_output: DataMapping,
            node_yields: Mapping[str, str],
        ):
            received["partial_output"] = partial_output
            return WorkflowExecutionResult.success(partial_output)

        context.on_workflow_yield = on_workflow_yield
        result = await algorithm.execute(
            context=context,
            workflow=workflow,
            input={},
        )

        assert result.status is WorkflowExecutionResultStatus.SUCCESS
        assert received["partial_output"] == {}
        assert result.output == {}


class TestWorkflowErrorNodeYields:
    @pytest.mark.asyncio
    async def test_node_yields_passed_to_on_workflow_error(self):
        """on_workflow_error receives node_yields (empty when no nodes yielded)."""
        workflow = _error_workflow()
        context = InMemoryContext()
        received: dict = {}

        async def on_workflow_error(
            *,
            workflow: Workflow,
            input: DataMapping,
            errors: WorkflowErrors,
            partial_output: DataMapping,
            node_yields: Mapping[str, str],
        ):
            received["node_yields"] = node_yields
            return WorkflowExecutionResult.error(
                errors=errors,
                partial_output=partial_output,
                node_yields=node_yields,
            )

        context.on_workflow_error = on_workflow_error

        algorithm = TopologicalExecutionAlgorithm()
        result = await algorithm.execute(context=context, workflow=workflow, input={})

        assert result.status is WorkflowExecutionResultStatus.ERROR
        assert received["node_yields"] == {}


class TestErrorPrecedenceOverYield:
    @pytest.mark.asyncio
    async def test_error_takes_precedence_over_yield_in_continue_mode(self):
        """When both errors and yields occur, on_workflow_error is called, not on_workflow_yield."""
        workflow = _error_and_yield_workflow()
        context = InMemoryContext()
        yield_mock = AsyncMock(side_effect=context.on_workflow_yield)
        error_mock = AsyncMock(side_effect=context.on_workflow_error)
        context.on_workflow_yield = yield_mock
        context.on_workflow_error = error_mock

        algorithm = ParallelExecutionAlgorithm(
            error_handling=ErrorHandlingMode.CONTINUE,
        )
        result = await algorithm.execute(context=context, workflow=workflow, input={})

        assert result.status is WorkflowExecutionResultStatus.ERROR
        yield_mock.assert_not_called()
        error_mock.assert_called_once()

    @pytest.mark.asyncio
    async def test_node_yields_included_in_error_hook_when_both_occur(self):
        """When both errors and yields occur, on_workflow_error receives the node_yields."""
        workflow = _error_and_yield_workflow()
        context = InMemoryContext()
        received: dict = {}

        async def on_workflow_error(
            *,
            workflow: Workflow,
            input: DataMapping,
            errors: WorkflowErrors,
            partial_output: DataMapping,
            node_yields: Mapping[str, str],
        ):
            received["node_yields"] = node_yields
            return WorkflowExecutionResult.error(
                errors=errors,
                partial_output=partial_output,
                node_yields=node_yields,
            )

        context.on_workflow_error = on_workflow_error

        algorithm = ParallelExecutionAlgorithm(
            error_handling=ErrorHandlingMode.CONTINUE,
        )
        result = await algorithm.execute(context=context, workflow=workflow, input={})

        assert result.status is WorkflowExecutionResultStatus.ERROR
        assert "yielding" in received["node_yields"]
