from __future__ import annotations

import asyncio
import sys
from types import ModuleType
from typing import Any

import pytest
from pydantic import BaseModel

from agentbyte.cancellation_token import CancellationToken
from agentbyte.workflow import (
    WorkflowApprovalRequestedEvent,
    WorkflowCancelledError,
    StepFailedEvent,
    StepMetadata,
    StepProgressEvent,
    StepStartedEvent,
    Workflow,
    WorkflowContext,
    WorkflowCancelledEvent,
    WorkflowCompletedEvent,
    WorkflowCompletionMode,
    WorkflowEventOrigin,
    WorkflowConfig,
    WorkflowFailedError,
    WorkflowFailedEvent,
    WorkflowInputRequestedEvent,
    WorkflowResumedEvent,
    WorkflowRunner,
    WorkflowStateConflictPolicy,
    WorkflowStateMode,
    WorkflowStatus,
    WorkflowSuspendedEvent,
)
from agentbyte.workflow.core.models import WorkflowContext as Context
from agentbyte.workflow.steps.echo import EchoOutput, EchoStep
from agentbyte.workflow.steps.function import FunctionStep
from agentbyte.workflow.steps.step import BaseStep


class NumberInput(BaseModel):
    value: int


class NumberOutput(BaseModel):
    result: int


def _build_math_workflow() -> Workflow:
    async def double(input_data: NumberInput, _: Context) -> NumberOutput:
        await asyncio.sleep(0)
        return NumberOutput(result=input_data.value * 2)

    async def add_ten(input_data: NumberOutput, _: Context) -> NumberOutput:
        await asyncio.sleep(0)
        return NumberOutput(result=input_data.result + 10)

    double_step = FunctionStep(
        step_id="double",
        metadata=StepMetadata(name="double"),
        input_type=NumberInput,
        output_type=NumberOutput,
        func=double,
    )
    add_ten_step = FunctionStep(
        step_id="add_ten",
        metadata=StepMetadata(name="add_ten"),
        input_type=NumberOutput,
        output_type=NumberOutput,
        func=add_ten,
    )
    return Workflow(WorkflowConfig(name="math-pipeline")).chain(double_step, add_ten_step)


@pytest.mark.asyncio
async def test_runner_supports_listing_6_2_style_fluent_chain_composition() -> None:
    class ListingInput(BaseModel):
        value: int

    class ListingOutput(BaseModel):
        result: int

    async def double(input_data: ListingInput, _: Context) -> ListingOutput:
        await asyncio.sleep(0)
        return ListingOutput(result=input_data.value * 2)

    async def add_ten(input_data: ListingOutput, _: Context) -> ListingOutput:
        await asyncio.sleep(0)
        return ListingOutput(result=input_data.result + 10)

    double_step = FunctionStep(
        step_id="double",
        metadata=StepMetadata(name="double"),
        input_type=ListingInput,
        output_type=ListingOutput,
        func=double,
    )
    add_ten_step = FunctionStep(
        step_id="add_ten",
        metadata=StepMetadata(name="add_ten"),
        input_type=ListingOutput,
        output_type=ListingOutput,
        func=add_ten,
    )

    workflow = (
        Workflow(WorkflowConfig(name="simple-math-pipeline"))
        .chain(double_step, add_ten_step)
    )

    assert workflow.start_step_id == "double"
    assert workflow.end_step_ids == ["add_ten"]

    runner = WorkflowRunner()
    execution = await runner.run(workflow, {"value": 7})

    assert execution.status.value == "completed"
    assert execution.step_executions["double"].output_data == {"result": 14}
    assert execution.step_executions["add_ten"].output_data == {"result": 24}


class _FakeStatusCode:
    OK = "OK"
    ERROR = "ERROR"


class _FakeStatus:
    def __init__(self, code: str, description: str | None = None) -> None:
        self.code = code
        self.description = description


class _FakeSpan:
    def __init__(self, name: str) -> None:
        self.name = name
        self.attributes: dict[str, Any] = {}
        self.status: _FakeStatus | None = None
        self.recorded_exceptions: list[Exception] = []
        self.parent_name: str | None = None

    def set_attribute(self, key: str, value: Any) -> None:
        self.attributes[key] = value

    def set_status(self, status: _FakeStatus) -> None:
        self.status = status

    def record_exception(self, exc: Exception) -> None:
        self.recorded_exceptions.append(exc)


class _FakeSpanContextManager:
    def __init__(self, span: _FakeSpan, active_stack: list[_FakeSpan]) -> None:
        self.span = span
        self._active_stack = active_stack

    def __enter__(self) -> _FakeSpan:
        if self._active_stack:
            self.span.parent_name = self._active_stack[-1].name
        self._active_stack.append(self.span)
        return self.span

    def __exit__(self, exc_type: Any, exc: Any, tb: Any) -> bool:
        self._active_stack.pop()
        return False


class _FakeTracer:
    def __init__(self, spans: list[_FakeSpan], active_stack: list[_FakeSpan]) -> None:
        self._spans = spans
        self._active_stack = active_stack

    def start_as_current_span(self, name: str) -> _FakeSpanContextManager:
        span = _FakeSpan(name)
        self._spans.append(span)
        return _FakeSpanContextManager(span, self._active_stack)


def _install_fake_opentelemetry(monkeypatch: pytest.MonkeyPatch, spans: list[_FakeSpan]) -> None:
    active_stack: list[_FakeSpan] = []
    trace_module = ModuleType("opentelemetry.trace")
    trace_module.get_tracer = lambda name: _FakeTracer(spans, active_stack)  # type: ignore[attr-defined]
    trace_module.Status = _FakeStatus  # type: ignore[attr-defined]
    trace_module.StatusCode = _FakeStatusCode  # type: ignore[attr-defined]

    otel_module = ModuleType("opentelemetry")
    otel_module.trace = trace_module  # type: ignore[attr-defined]

    monkeypatch.setitem(sys.modules, "opentelemetry", otel_module)
    monkeypatch.setitem(sys.modules, "opentelemetry.trace", trace_module)


@pytest.mark.asyncio
async def test_runner_completes_sequential_workflow() -> None:
    workflow = _build_math_workflow()
    runner = WorkflowRunner()

    execution = await runner.run(workflow, {"value": 3})

    assert execution.status.value == "completed"
    assert execution.step_executions["double"].output_data == {"result": 6}
    assert execution.step_executions["add_ten"].output_data == {"result": 16}
    assert execution.state["double_output"] == {"result": 6}
    assert execution.state["add_ten_output"] == {"result": 16}


@pytest.mark.asyncio
async def test_runner_emits_progress_events() -> None:
    class ProgressStep(BaseStep[NumberInput, NumberOutput]):
        def __init__(self) -> None:
            super().__init__(
                step_id="progress",
                metadata=StepMetadata(name="progress"),
                input_type=NumberInput,
                output_type=NumberOutput,
            )

        async def execute(self, input_data: NumberInput, context: Context) -> NumberOutput:
            context.emit_progress("working", completed=1, total=2)
            await asyncio.sleep(0)
            context.emit_progress("done", completed=2, total=2)
            return NumberOutput(result=input_data.value)

    workflow = Workflow(WorkflowConfig(name="progress-flow"))
    workflow.set_start_step(ProgressStep()).add_end_step("progress")
    runner = WorkflowRunner()

    events = [event async for event in runner.run_stream(workflow, {"value": 5})]

    progress_events = [event for event in events if isinstance(event, StepProgressEvent)]
    assert len(progress_events) == 2
    assert progress_events[0].message == "working"
    assert isinstance(events[-1], WorkflowCompletedEvent)


@pytest.mark.asyncio
async def test_runner_completed_event_exposes_final_output_from_end_step() -> None:
    workflow = _build_math_workflow()
    runner = WorkflowRunner()

    events = [event async for event in runner.run_stream(workflow, {"value": 3})]

    completed_event = next(
        event for event in events if isinstance(event, WorkflowCompletedEvent)
    )

    assert completed_event.final_output == {"result": 16}
    assert completed_event.execution.step_executions["add_ten"].output_data == {"result": 16}


@pytest.mark.asyncio
async def test_runner_completes_when_end_step_reached_in_conditional_flow() -> None:
    start = EchoStep(step_id="start", metadata=StepMetadata(name="start"))
    go = EchoStep(
        step_id="go",
        metadata=StepMetadata(name="go"),
        input_type=EchoOutput,
        output_type=EchoOutput,
    )
    stop = EchoStep(
        step_id="stop",
        metadata=StepMetadata(name="stop"),
        input_type=EchoOutput,
        output_type=EchoOutput,
    )

    workflow = Workflow(WorkflowConfig(name="branch-flow"))
    workflow.add_step(start).add_step(go).add_step(stop)
    workflow.set_start_step(start).add_end_step(go)
    workflow.add_edge(
        start,
        go,
        {"type": "output_based", "field": "result", "operator": "==", "value": "go"},
    )
    workflow.add_edge(
        start,
        stop,
        {"type": "output_based", "field": "result", "operator": "==", "value": "stop"},
    )

    runner = WorkflowRunner()
    execution = await runner.run(workflow, {"message": "go"})

    assert execution.status.value == "completed"
    assert "go" in execution.step_executions
    assert "stop" not in execution.step_executions


@pytest.mark.asyncio
async def test_runner_detects_deadlock_when_no_steps_can_progress() -> None:
    start = EchoStep(step_id="start", metadata=StepMetadata(name="start"))
    blocked = EchoStep(
        step_id="blocked",
        metadata=StepMetadata(name="blocked"),
        input_type=EchoOutput,
        output_type=EchoOutput,
    )
    workflow = Workflow(WorkflowConfig(name="deadlock-flow"))
    workflow.add_edge(
        start,
        blocked,
        {"type": "output_based", "field": "result", "operator": "==", "value": "never"},
    ).set_start_step(start)

    runner = WorkflowRunner()

    with pytest.raises(RuntimeError, match="cannot be executed"):
        await runner.run(workflow, {"message": "hello"})


@pytest.mark.asyncio
async def test_runner_suspends_and_resumes_with_coerced_response() -> None:
    class DraftInput(BaseModel):
        draft: str

    class ReviewDecision(BaseModel):
        approved: bool
        comment: str

    class DraftOutput(BaseModel):
        result: str

    class HumanReviewStep(BaseStep[DraftInput, DraftOutput]):
        def __init__(self) -> None:
            super().__init__(
                step_id="review",
                metadata=StepMetadata(name="review"),
                input_type=DraftInput,
                output_type=DraftOutput,
            )

        async def execute(
            self, input_data: DraftInput, context: WorkflowContext
        ) -> DraftOutput:
            await context.request_input(
                key="review_feedback",
                response_type=ReviewDecision,
                prompt=f"Review draft: {input_data.draft}",
            )
            raise AssertionError("request_input should suspend execution")

        async def resume_from_request(
            self,
            input_data: DraftInput,
            context: WorkflowContext,
            request,
            response: ReviewDecision,
        ) -> DraftOutput:
            context.set("review_comment", response.comment)
            status = "approved" if response.approved else "rejected"
            return DraftOutput(result=f"{input_data.draft}:{status}:{response.comment}")

    workflow = Workflow(WorkflowConfig(name="human-review-flow"))
    workflow.set_start_step(HumanReviewStep()).add_end_step("review")
    runner = WorkflowRunner()

    events = [event async for event in runner.run_stream(workflow, {"draft": "spec-01"})]

    request_events = [
        event for event in events if isinstance(event, WorkflowInputRequestedEvent)
    ]
    assert len(request_events) == 1
    request_event = request_events[0]
    assert request_event.step_id == "review"
    assert request_event.request.key == "review_feedback"
    assert request_event.request.response_type_name == "ReviewDecision"

    suspended_events = [
        event for event in events if isinstance(event, WorkflowSuspendedEvent)
    ]
    assert len(suspended_events) == 1
    assert suspended_events[0].execution.status == WorkflowStatus.SUSPENDED

    resumed_events = [
        event
        async for event in runner.run_stream(
            workflow,
            responses={
                request_event.request.request_id: {
                    "approved": True,
                    "comment": "ship it",
                }
            },
        )
    ]
    execution = next(
        event.execution
        for event in resumed_events
        if isinstance(event, WorkflowCompletedEvent)
    )

    resumed_markers = [
        event for event in resumed_events if isinstance(event, WorkflowResumedEvent)
    ]
    assert len(resumed_markers) == 1
    resumed_step_started = [
        event
        for event in resumed_events
        if isinstance(event, StepStartedEvent) and event.step_id == "review"
    ]
    assert len(resumed_step_started) == 1

    assert execution.status == WorkflowStatus.COMPLETED
    assert execution.step_executions["review"].output_data == {
        "result": "spec-01:approved:ship it"
    }
    assert execution.state["review_comment"] == "ship it"
    assert execution.pending_requests == {}


@pytest.mark.asyncio
async def test_runner_emits_approval_request_event() -> None:
    class PublishInput(BaseModel):
        title: str

    class PublishOutput(BaseModel):
        result: str

    class PublishStep(BaseStep[PublishInput, PublishOutput]):
        def __init__(self) -> None:
            super().__init__(
                step_id="publish",
                metadata=StepMetadata(name="publish"),
                input_type=PublishInput,
                output_type=PublishOutput,
            )

        async def execute(
            self, input_data: PublishInput, context: WorkflowContext
        ) -> PublishOutput:
            await context.request_approval(
                key="publish_report",
                reason=f"Approve publish for {input_data.title}?",
                payload={"title": input_data.title},
            )
            raise AssertionError("request_approval should suspend execution")

        async def resume_from_request(
            self,
            input_data: PublishInput,
            context: WorkflowContext,
            request,
            response: bool,
        ) -> PublishOutput:
            context.set("approved", response)
            status = "approved" if response else "rejected"
            return PublishOutput(result=f"{input_data.title}:{status}:{request.key}")

    workflow = Workflow(WorkflowConfig(name="approval-flow"))
    workflow.set_start_step(PublishStep()).add_end_step("publish")
    runner = WorkflowRunner()

    events = [event async for event in runner.run_stream(workflow, {"title": "Guide"})]

    approval_events = [
        event for event in events if isinstance(event, WorkflowApprovalRequestedEvent)
    ]
    assert len(approval_events) == 1
    assert approval_events[0].request.key == "publish_report"
    assert approval_events[0].request.kind.value == "approval"


@pytest.mark.asyncio
async def test_runner_invalid_resume_response_fails_validation() -> None:
    class DraftInput(BaseModel):
        draft: str

    class ReviewDecision(BaseModel):
        approved: bool
        comment: str

    class DraftOutput(BaseModel):
        result: str

    class HumanReviewStep(BaseStep[DraftInput, DraftOutput]):
        def __init__(self) -> None:
            super().__init__(
                step_id="review",
                metadata=StepMetadata(name="review"),
                input_type=DraftInput,
                output_type=DraftOutput,
            )

        async def execute(
            self, input_data: DraftInput, context: WorkflowContext
        ) -> DraftOutput:
            await context.request_input(
                key="review_feedback",
                response_type=ReviewDecision,
                prompt=f"Review draft: {input_data.draft}",
            )
            raise AssertionError("request_input should suspend execution")

    workflow = Workflow(WorkflowConfig(name="invalid-response-flow"))
    workflow.set_start_step(HumanReviewStep()).add_end_step("review")
    runner = WorkflowRunner()

    events = [event async for event in runner.run_stream(workflow, {"draft": "spec-01"})]
    request_event = next(
        event for event in events if isinstance(event, WorkflowInputRequestedEvent)
    )

    with pytest.raises(WorkflowFailedError, match="comment"):
        await runner.run(
            workflow,
            responses={
                request_event.request.request_id: {
                    "approved": True,
                }
            },
        )


@pytest.mark.asyncio
async def test_runner_tracks_and_resolves_multiple_pending_requests() -> None:
    class StartInput(BaseModel):
        value: int

    class StartOutput(BaseModel):
        result: int

    class ReviewDecision(BaseModel):
        approved: bool
        note: str

    class ReviewOutput(BaseModel):
        result: str

    async def start_step(input_data: StartInput, _: Context) -> StartOutput:
        return StartOutput(result=input_data.value)

    class SuspendStep(BaseStep[StartOutput, ReviewOutput]):
        def __init__(self, step_id: str) -> None:
            super().__init__(
                step_id=step_id,
                metadata=StepMetadata(name=step_id),
                input_type=StartOutput,
                output_type=ReviewOutput,
            )

        async def execute(
            self, input_data: StartOutput, context: WorkflowContext
        ) -> ReviewOutput:
            await context.request_input(
                key=f"{self.step_id}_feedback",
                response_type=ReviewDecision,
                prompt=f"Provide feedback for {self.step_id}:{input_data.result}",
            )
            raise AssertionError("request_input should suspend execution")

        async def resume_from_request(
            self,
            input_data: StartOutput,
            context: WorkflowContext,
            request,
            response: ReviewDecision,
        ) -> ReviewOutput:
            context.set(request.key, response.note)
            return ReviewOutput(result=f"{self.step_id}:{response.note}")

    start = FunctionStep(
        step_id="start",
        metadata=StepMetadata(name="start"),
        input_type=StartInput,
        output_type=StartOutput,
        func=start_step,
    )
    left = SuspendStep("left")
    right = SuspendStep("right")

    workflow = Workflow(
        WorkflowConfig(
            name="multi-request-flow",
            completion_mode=WorkflowCompletionMode.ALL_END,
        )
    )
    workflow.add_step(start).add_step(left).add_step(right)
    workflow.set_start_step(start).add_end_step(left).add_end_step(right)
    workflow.add_edge(start, left)
    workflow.add_edge(start, right)
    runner = WorkflowRunner(max_concurrent_steps=3)

    initial_events = [
        event async for event in runner.run_stream(workflow, {"value": 5})
    ]
    request_events = [
        event for event in initial_events if isinstance(event, WorkflowInputRequestedEvent)
    ]
    assert len(request_events) == 2
    request_by_step = {event.step_id: event for event in request_events}
    suspended_event = next(
        event for event in initial_events if isinstance(event, WorkflowSuspendedEvent)
    )
    assert len(suspended_event.execution.pending_requests) == 2

    resumed_execution = await runner.run(
        workflow,
        responses={
            request_by_step["left"].request.request_id: {
                "approved": True,
                "note": "left-ok",
            },
            request_by_step["right"].request.request_id: {
                "approved": False,
                "note": "right-ok",
            },
        },
    )

    assert resumed_execution.status == WorkflowStatus.COMPLETED
    assert resumed_execution.pending_requests == {}
    assert resumed_execution.step_executions["left"].output_data == {
        "result": "left:left-ok"
    }
    assert resumed_execution.step_executions["right"].output_data == {
        "result": "right:right-ok"
    }
    assert resumed_execution.state["left_feedback"] == "left-ok"
    assert resumed_execution.state["right_feedback"] == "right-ok"


@pytest.mark.asyncio
async def test_runner_rejects_reentrant_execution_for_same_workflow() -> None:
    started = asyncio.Event()
    unblock = asyncio.Event()

    class BlockingStep(BaseStep[NumberInput, NumberOutput]):
        def __init__(self) -> None:
            super().__init__(
                step_id="block",
                metadata=StepMetadata(name="block"),
                input_type=NumberInput,
                output_type=NumberOutput,
            )

        async def execute(
            self, input_data: NumberInput, context: WorkflowContext
        ) -> NumberOutput:
            started.set()
            await unblock.wait()
            context.set("done", True)
            return NumberOutput(result=input_data.value)

    workflow = Workflow(WorkflowConfig(name="reentrant-flow"))
    workflow.set_start_step(BlockingStep()).add_end_step("block")
    runner = WorkflowRunner()

    task = asyncio.create_task(runner.run(workflow, {"value": 1}))
    await started.wait()

    with pytest.raises(RuntimeError, match="already running"):
        await runner.run(workflow, {"value": 2})

    unblock.set()
    execution = await task
    assert execution.status == WorkflowStatus.COMPLETED


@pytest.mark.asyncio
async def test_runner_run_raises_distinct_cancellation_error() -> None:
    class SlowStep(BaseStep[NumberInput, NumberOutput]):
        def __init__(self) -> None:
            super().__init__(
                step_id="slow",
                metadata=StepMetadata(name="slow"),
                input_type=NumberInput,
                output_type=NumberOutput,
            )

        async def execute(self, input_data: NumberInput, context: Context) -> NumberOutput:  # noqa: ARG002
            await asyncio.sleep(0.2)
            return NumberOutput(result=input_data.value)

    workflow = Workflow(WorkflowConfig(name="cancel-error-flow"))
    workflow.set_start_step(SlowStep()).add_end_step("slow")
    runner = WorkflowRunner()
    token = CancellationToken()

    async def cancel_soon() -> None:
        await asyncio.sleep(0.01)
        token.cancel()

    cancel_task = asyncio.create_task(cancel_soon())
    try:
        with pytest.raises(WorkflowCancelledError, match="Cancelled by user"):
            await runner.run(workflow, {"value": 1}, cancellation_token=token)
    finally:
        await cancel_task


@pytest.mark.asyncio
async def test_runner_supports_cancellation() -> None:
    class SlowStep(BaseStep[NumberInput, NumberOutput]):
        def __init__(self) -> None:
            super().__init__(
                step_id="slow",
                metadata=StepMetadata(name="slow"),
                input_type=NumberInput,
                output_type=NumberOutput,
            )

        async def execute(self, input_data: NumberInput, context: Context) -> NumberOutput:  # noqa: ARG002
            await asyncio.sleep(0.2)
            return NumberOutput(result=input_data.value)

    workflow = Workflow(WorkflowConfig(name="cancel-flow"))
    workflow.set_start_step(SlowStep()).add_end_step("slow")
    runner = WorkflowRunner()
    token = CancellationToken()

    events: list[Any] = []
    async for event in runner.run_stream(workflow, {"value": 1}, token):
        events.append(event)
        if getattr(event, "step_id", None) == "slow":
            token.cancel()

    assert any(isinstance(event, WorkflowCancelledEvent) for event in events)


@pytest.mark.asyncio
async def test_runner_creates_workflow_and_step_spans_when_otel_enabled(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    workflow = _build_math_workflow()
    runner = WorkflowRunner()
    spans: list[_FakeSpan] = []

    monkeypatch.setenv("AGENTBYTE_ENABLE_OTEL", "true")
    _install_fake_opentelemetry(monkeypatch, spans)

    execution = await runner.run(workflow, {"value": 3})

    assert execution.status.value == "completed"
    assert [span.name for span in spans] == [
        "workflow math-pipeline",
        "workflow step double",
        "workflow step add_ten",
    ]

    workflow_span = spans[0]
    assert workflow_span.attributes["agentbyte.workflow.name"] == "math-pipeline"
    assert workflow_span.attributes["agentbyte.workflow.status"] == "completed"
    assert workflow_span.attributes["agentbyte.workflow.completed_steps"] == 2
    assert workflow_span.status is not None
    assert workflow_span.status.code == _FakeStatusCode.OK

    first_step_span = spans[1]
    assert first_step_span.attributes["agentbyte.workflow.step.id"] == "double"
    assert first_step_span.attributes["agentbyte.workflow.step.status"] == "completed"
    assert first_step_span.attributes["agentbyte.workflow.step.output_keys"] == "result"
    assert first_step_span.parent_name == "workflow math-pipeline"


@pytest.mark.asyncio
async def test_runner_preserves_nested_span_hierarchy_inside_step(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    class InnerSpanStep(BaseStep[NumberInput, NumberOutput]):
        def __init__(self) -> None:
            super().__init__(
                step_id="inner",
                metadata=StepMetadata(name="inner"),
                input_type=NumberInput,
                output_type=NumberOutput,
            )

        async def execute(self, input_data: NumberInput, context: Context) -> NumberOutput:  # noqa: ARG002
            from opentelemetry import trace

            tracer = trace.get_tracer("agentbyte.test")
            with tracer.start_as_current_span("inner operation"):
                await asyncio.sleep(0)
                return NumberOutput(result=input_data.value)

    workflow = Workflow(WorkflowConfig(name="nested-span-flow"))
    workflow.set_start_step(InnerSpanStep()).add_end_step("inner")
    runner = WorkflowRunner()
    spans: list[_FakeSpan] = []

    monkeypatch.setenv("AGENTBYTE_ENABLE_OTEL", "true")
    _install_fake_opentelemetry(monkeypatch, spans)

    execution = await runner.run(workflow, {"value": 7})

    assert execution.status.value == "completed"
    assert [span.name for span in spans] == [
        "workflow nested-span-flow",
        "workflow step inner",
        "inner operation",
    ]
    assert spans[1].parent_name == "workflow nested-span-flow"
    assert spans[2].parent_name == "workflow step inner"


@pytest.mark.asyncio
async def test_runner_marks_step_failure_on_step_span_without_workflow_status(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    class FailingStep(BaseStep[NumberInput, NumberOutput]):
        def __init__(self) -> None:
            super().__init__(
                step_id="fail",
                metadata=StepMetadata(name="fail"),
                input_type=NumberInput,
                output_type=NumberOutput,
            )

        async def execute(
            self,
            input_data: NumberInput,
            context: Context,
        ) -> NumberOutput:
            del input_data, context
            raise RuntimeError("boom")

    workflow = Workflow(WorkflowConfig(name="step-failure-span-flow"))
    workflow.set_start_step(FailingStep()).add_end_step("fail")
    runner = WorkflowRunner()
    spans: list[_FakeSpan] = []

    monkeypatch.setenv("AGENTBYTE_ENABLE_OTEL", "true")
    _install_fake_opentelemetry(monkeypatch, spans)

    events = [event async for event in runner.run_stream(workflow, {"value": 7})]
    assert isinstance(events[-1], WorkflowFailedEvent)

    workflow_span = spans[0]
    step_span = spans[1]

    assert workflow_span.attributes["agentbyte.workflow.status"] == "failed"
    assert step_span.attributes["agentbyte.workflow.step.status"] == "failed"
    assert "agentbyte.workflow.status" not in step_span.attributes
    assert step_span.attributes["agentbyte.workflow.step.type"] == "FailingStep"
    assert step_span.recorded_exceptions
    assert step_span.status is not None
    assert step_span.status.code == _FakeStatusCode.ERROR


@pytest.mark.asyncio
async def test_runner_surfaces_input_validation_failure() -> None:
    workflow = _build_math_workflow()
    runner = WorkflowRunner()

    events = [event async for event in runner.run_stream(workflow, {"bad": 3})]

    assert isinstance(events[-1], WorkflowFailedEvent)
    assert "Initial input validation failed" in events[-1].error
    assert events[-1].origin is WorkflowEventOrigin.FRAMEWORK
    assert events[-1].error_details.error_type == "ValidationError"
    assert events[-1].error_details.step_id == workflow.start_step_id
    assert events[-1].error_details.workflow_id == workflow.id


@pytest.mark.asyncio
async def test_runner_emits_structured_step_failure_details_with_traceback() -> None:
    async def explode(input_data: NumberInput, _: WorkflowContext) -> NumberOutput:
        del input_data
        raise RuntimeError("boom")

    workflow = Workflow(WorkflowConfig(name="runtime-step-failure"))
    workflow.set_start_step(
        FunctionStep(
            step_id="explode",
            metadata=StepMetadata(name="explode"),
            input_type=NumberInput,
            output_type=NumberOutput,
            func=explode,
        )
    ).add_end_step("explode")

    events = [event async for event in WorkflowRunner().run_stream(workflow, {"value": 3})]

    step_failed = next(event for event in events if isinstance(event, StepFailedEvent))
    workflow_failed = events[-1]

    assert isinstance(workflow_failed, WorkflowFailedEvent)
    assert step_failed.step_id == "explode"
    assert step_failed.error_details.error_type == "RuntimeError"
    assert step_failed.error_details.message == "boom"
    assert step_failed.error_details.step_id == "explode"
    assert step_failed.error_details.workflow_id == workflow.id
    assert step_failed.error_details.execution_id is not None
    assert step_failed.error_details.traceback is not None
    assert "RuntimeError: boom" in step_failed.error_details.traceback
    assert workflow_failed.error_details.traceback is not None
    assert workflow_failed.error_details.execution_id == step_failed.error_details.execution_id


@pytest.mark.asyncio
async def test_runner_collects_all_fan_in_branch_outputs_as_envelope() -> None:
    """Fan-in join step must receive branch_outputs envelope with outputs from all branches."""

    class BranchOutput(BaseModel):
        value: int

    class FanInInput(BaseModel):
        branch_outputs: dict[str, Any]
        failed_branches: list[str] = []

    received: list[FanInInput] = []

    async def branch_a_fn(input_data: NumberInput, _: Context) -> BranchOutput:
        return BranchOutput(value=input_data.value * 2)

    async def branch_b_fn(input_data: NumberInput, _: Context) -> BranchOutput:
        return BranchOutput(value=input_data.value + 100)

    async def join_fn(input_data: FanInInput, _: Context) -> NumberOutput:
        received.append(input_data)
        total = sum(v["value"] for v in input_data.branch_outputs.values())
        return NumberOutput(result=total)

    start = FunctionStep(
        step_id="start",
        metadata=StepMetadata(name="start"),
        input_type=NumberInput,
        output_type=NumberInput,
        func=lambda i, c: i,
    )
    branch_a = FunctionStep(
        step_id="branch_a",
        metadata=StepMetadata(name="branch_a"),
        input_type=NumberInput,
        output_type=BranchOutput,
        func=branch_a_fn,
    )
    branch_b = FunctionStep(
        step_id="branch_b",
        metadata=StepMetadata(name="branch_b"),
        input_type=NumberInput,
        output_type=BranchOutput,
        func=branch_b_fn,
    )
    join = FunctionStep(
        step_id="join",
        metadata=StepMetadata(name="join"),
        input_type=FanInInput,
        output_type=NumberOutput,
        func=join_fn,
    )

    workflow = Workflow(WorkflowConfig(name="fan-in-flow"))
    workflow.set_start_step(start)
    workflow.add_edge(start, branch_a)
    workflow.add_edge(start, branch_b)
    workflow.add_edge(branch_a, join)
    workflow.add_edge(branch_b, join)
    workflow.add_end_step(join)

    runner = WorkflowRunner()
    execution = await runner.run(workflow, {"value": 5})

    assert execution.status.value == "completed"
    assert len(received) == 1
    env = received[0]
    assert env.failed_branches == []
    assert set(env.branch_outputs.keys()) == {"branch_a", "branch_b"}
    assert env.branch_outputs["branch_a"]["value"] == 10   # 5 * 2
    assert env.branch_outputs["branch_b"]["value"] == 105  # 5 + 100


@pytest.mark.asyncio
async def test_runner_waits_for_all_end_steps_in_all_end_mode() -> None:
    """With ALL_END mode the runner must wait for ALL declared end steps to complete."""

    async def left_fn(input_data: NumberInput, _: Context) -> NumberOutput:
        return NumberOutput(result=input_data.value + 1)

    async def right_fn(input_data: NumberInput, _: Context) -> NumberOutput:
        return NumberOutput(result=input_data.value + 2)

    start = FunctionStep(
        step_id="start",
        metadata=StepMetadata(name="start"),
        input_type=NumberInput,
        output_type=NumberInput,
        func=lambda i, c: i,
    )
    left = FunctionStep(
        step_id="left",
        metadata=StepMetadata(name="left"),
        input_type=NumberInput,
        output_type=NumberOutput,
        func=left_fn,
    )
    right = FunctionStep(
        step_id="right",
        metadata=StepMetadata(name="right"),
        input_type=NumberInput,
        output_type=NumberOutput,
        func=right_fn,
    )

    workflow = Workflow(WorkflowConfig(name="all-end-flow", completion_mode=WorkflowCompletionMode.ALL_END))
    workflow.set_start_step(start)
    workflow.add_edge(start, left)
    workflow.add_edge(start, right)
    workflow.add_end_step(left)
    workflow.add_end_step(right)

    runner = WorkflowRunner()
    execution = await runner.run(workflow, {"value": 10})

    assert execution.status.value == "completed"
    # Both declared end steps must have completed
    assert execution.step_executions["left"].status.value == "completed"
    assert execution.step_executions["right"].status.value == "completed"
    assert execution.step_executions["left"].output_data == {"result": 11}
    assert execution.step_executions["right"].output_data == {"result": 12}


@pytest.mark.asyncio
async def test_runner_immediate_state_mode_preserves_parallel_state_visibility() -> None:
    class FanInInput(BaseModel):
        branch_outputs: dict[str, Any]
        failed_branches: list[str] = []

    class ObservationOutput(BaseModel):
        saw_shared: bool

    start_signal = asyncio.Event()

    async def start_fn(input_data: NumberInput, _: Context) -> NumberInput:
        return input_data

    async def branch_a_fn(input_data: NumberInput, context: WorkflowContext) -> NumberOutput:
        del input_data
        context.set("shared", "visible")
        start_signal.set()
        await asyncio.sleep(0)
        return NumberOutput(result=1)

    async def branch_b_fn(
        input_data: NumberInput,
        context: WorkflowContext,
    ) -> ObservationOutput:
        del input_data
        await start_signal.wait()
        return ObservationOutput(saw_shared=context.get("shared") == "visible")

    async def join_fn(input_data: FanInInput, _: Context) -> NumberOutput:
        saw_shared = input_data.branch_outputs["branch_b"]["saw_shared"]
        return NumberOutput(result=1 if saw_shared else 0)

    start = FunctionStep(
        step_id="start",
        metadata=StepMetadata(name="start"),
        input_type=NumberInput,
        output_type=NumberInput,
        func=start_fn,
    )
    branch_a = FunctionStep(
        step_id="branch_a",
        metadata=StepMetadata(name="branch_a"),
        input_type=NumberInput,
        output_type=NumberOutput,
        func=branch_a_fn,
    )
    branch_b = FunctionStep(
        step_id="branch_b",
        metadata=StepMetadata(name="branch_b"),
        input_type=NumberInput,
        output_type=ObservationOutput,
        func=branch_b_fn,
    )
    join = FunctionStep(
        step_id="join",
        metadata=StepMetadata(name="join"),
        input_type=FanInInput,
        output_type=NumberOutput,
        func=join_fn,
    )

    workflow = Workflow(WorkflowConfig(name="immediate-state-visibility"))
    workflow.set_start_step(start)
    workflow.add_edge(start, branch_a)
    workflow.add_edge(start, branch_b)
    workflow.add_edge(branch_a, join)
    workflow.add_edge(branch_b, join)
    workflow.add_end_step(join)

    execution = await WorkflowRunner().run(workflow, {"value": 1})

    assert execution.status == WorkflowStatus.COMPLETED
    assert execution.step_executions["join"].output_data == {"result": 1}
    assert execution.state["shared"] == "visible"


@pytest.mark.asyncio
async def test_runner_staged_state_mode_isolates_parallel_writes_until_commit() -> None:
    class FanInInput(BaseModel):
        branch_outputs: dict[str, Any]
        failed_branches: list[str] = []

    class ObservationOutput(BaseModel):
        saw_shared: bool

    start_signal = asyncio.Event()

    async def start_fn(input_data: NumberInput, _: Context) -> NumberInput:
        return input_data

    async def branch_a_fn(input_data: NumberInput, context: WorkflowContext) -> NumberOutput:
        del input_data
        context.set("shared", "visible")
        start_signal.set()
        await asyncio.sleep(0)
        return NumberOutput(result=1)

    async def branch_b_fn(
        input_data: NumberInput,
        context: WorkflowContext,
    ) -> ObservationOutput:
        del input_data
        await start_signal.wait()
        return ObservationOutput(saw_shared=context.get("shared") == "visible")

    async def join_fn(input_data: FanInInput, context: WorkflowContext) -> NumberOutput:
        saw_shared = input_data.branch_outputs["branch_b"]["saw_shared"]
        committed_shared = context.get("shared")
        return NumberOutput(result=2 if saw_shared else 1 if committed_shared == "visible" else 0)

    start = FunctionStep(
        step_id="start",
        metadata=StepMetadata(name="start"),
        input_type=NumberInput,
        output_type=NumberInput,
        func=start_fn,
    )
    branch_a = FunctionStep(
        step_id="branch_a",
        metadata=StepMetadata(name="branch_a"),
        input_type=NumberInput,
        output_type=NumberOutput,
        func=branch_a_fn,
    )
    branch_b = FunctionStep(
        step_id="branch_b",
        metadata=StepMetadata(name="branch_b"),
        input_type=NumberInput,
        output_type=ObservationOutput,
        func=branch_b_fn,
    )
    join = FunctionStep(
        step_id="join",
        metadata=StepMetadata(name="join"),
        input_type=FanInInput,
        output_type=NumberOutput,
        func=join_fn,
    )

    workflow = Workflow(
        WorkflowConfig(
            name="staged-state-visibility",
            state_mode=WorkflowStateMode.STAGED,
        )
    )
    workflow.set_start_step(start)
    workflow.add_edge(start, branch_a)
    workflow.add_edge(start, branch_b)
    workflow.add_edge(branch_a, join)
    workflow.add_edge(branch_b, join)
    workflow.add_end_step(join)

    execution = await WorkflowRunner().run(workflow, {"value": 1})

    assert execution.status == WorkflowStatus.COMPLETED
    assert execution.step_executions["branch_b"].output_data == {"saw_shared": False}
    assert execution.step_executions["join"].output_data == {"result": 1}
    assert execution.state["shared"] == "visible"


@pytest.mark.asyncio
async def test_runner_staged_state_mode_commits_parallel_branch_writes_before_join() -> None:
    class FanInInput(BaseModel):
        branch_outputs: dict[str, Any]
        failed_branches: list[str] = []

    async def start_fn(input_data: NumberInput, _: Context) -> NumberInput:
        return input_data

    async def left_fn(input_data: NumberInput, context: WorkflowContext) -> NumberOutput:
        del input_data
        context.set("left_total", 1)
        return NumberOutput(result=1)

    async def right_fn(input_data: NumberInput, context: WorkflowContext) -> NumberOutput:
        del input_data
        context.set("right_total", 2)
        return NumberOutput(result=2)

    async def join_fn(input_data: FanInInput, context: WorkflowContext) -> NumberOutput:
        del input_data
        return NumberOutput(
            result=context.get("left_total", 0) + context.get("right_total", 0)
        )

    start = FunctionStep(
        step_id="start",
        metadata=StepMetadata(name="start"),
        input_type=NumberInput,
        output_type=NumberInput,
        func=start_fn,
    )
    left = FunctionStep(
        step_id="left",
        metadata=StepMetadata(name="left"),
        input_type=NumberInput,
        output_type=NumberOutput,
        func=left_fn,
    )
    right = FunctionStep(
        step_id="right",
        metadata=StepMetadata(name="right"),
        input_type=NumberInput,
        output_type=NumberOutput,
        func=right_fn,
    )
    join = FunctionStep(
        step_id="join",
        metadata=StepMetadata(name="join"),
        input_type=FanInInput,
        output_type=NumberOutput,
        func=join_fn,
    )

    workflow = Workflow(
        WorkflowConfig(
            name="staged-state-join-commit",
            state_mode=WorkflowStateMode.STAGED,
        )
    )
    workflow.set_start_step(start)
    workflow.add_edge(start, left)
    workflow.add_edge(start, right)
    workflow.add_edge(left, join)
    workflow.add_edge(right, join)
    workflow.add_end_step(join)

    execution = await WorkflowRunner().run(workflow, {"value": 1})

    assert execution.status == WorkflowStatus.COMPLETED
    assert execution.step_executions["join"].output_data == {"result": 3}
    assert execution.state["left_total"] == 1
    assert execution.state["right_total"] == 2


@pytest.mark.asyncio
async def test_runner_staged_state_mode_fails_on_parallel_write_conflicts() -> None:
    async def start_fn(input_data: NumberInput, _: Context) -> NumberInput:
        return input_data

    async def left_fn(input_data: NumberInput, context: WorkflowContext) -> NumberOutput:
        del input_data
        context.set("shared", "left")
        return NumberOutput(result=1)

    async def right_fn(input_data: NumberInput, context: WorkflowContext) -> NumberOutput:
        del input_data
        context.set("shared", "right")
        return NumberOutput(result=2)

    start = FunctionStep(
        step_id="start",
        metadata=StepMetadata(name="start"),
        input_type=NumberInput,
        output_type=NumberInput,
        func=start_fn,
    )
    left = FunctionStep(
        step_id="left",
        metadata=StepMetadata(name="left"),
        input_type=NumberInput,
        output_type=NumberOutput,
        func=left_fn,
    )
    right = FunctionStep(
        step_id="right",
        metadata=StepMetadata(name="right"),
        input_type=NumberInput,
        output_type=NumberOutput,
        func=right_fn,
    )

    workflow = Workflow(
        WorkflowConfig(
            name="staged-state-conflict",
            state_mode=WorkflowStateMode.STAGED,
            state_conflict_policy=WorkflowStateConflictPolicy.FAIL,
        )
    )
    workflow.set_start_step(start)
    workflow.add_edge(start, left)
    workflow.add_edge(start, right)
    workflow.add_end_step(left)
    workflow.add_end_step(right)

    with pytest.raises(WorkflowFailedError, match="Staged workflow state conflict"):
        await WorkflowRunner().run(workflow, {"value": 1})


@pytest.mark.asyncio
async def test_runner_staged_state_mode_last_write_wins_uses_completion_order() -> None:
    zeta_finished = asyncio.Event()

    async def start_fn(input_data: NumberInput, _: Context) -> NumberInput:
        return input_data

    async def alpha_fn(input_data: NumberInput, context: WorkflowContext) -> NumberOutput:
        del input_data
        await zeta_finished.wait()
        await asyncio.sleep(0)
        context.set("shared", "alpha")
        return NumberOutput(result=1)

    async def zeta_fn(input_data: NumberInput, context: WorkflowContext) -> NumberOutput:
        del input_data
        context.set("shared", "zeta")
        zeta_finished.set()
        return NumberOutput(result=2)

    start = FunctionStep(
        step_id="start",
        metadata=StepMetadata(name="start"),
        input_type=NumberInput,
        output_type=NumberInput,
        func=start_fn,
    )
    alpha = FunctionStep(
        step_id="alpha",
        metadata=StepMetadata(name="alpha"),
        input_type=NumberInput,
        output_type=NumberOutput,
        func=alpha_fn,
    )
    zeta = FunctionStep(
        step_id="zeta",
        metadata=StepMetadata(name="zeta"),
        input_type=NumberInput,
        output_type=NumberOutput,
        func=zeta_fn,
    )

    workflow = Workflow(
        WorkflowConfig(
            name="staged-state-last-write-wins",
            state_mode=WorkflowStateMode.STAGED,
            state_conflict_policy=WorkflowStateConflictPolicy.LAST_WRITE_WINS,
        )
    )
    workflow.set_start_step(start)
    workflow.add_edge(start, alpha)
    workflow.add_edge(start, zeta)
    workflow.add_end_step(alpha)
    workflow.add_end_step(zeta)

    execution = await WorkflowRunner().run(workflow, {"value": 1})

    assert execution.status == WorkflowStatus.COMPLETED
    assert execution.state["shared"] == "alpha"


@pytest.mark.asyncio
async def test_runner_staged_state_mode_supports_state_deletion() -> None:
    async def delete_flag(input_data: NumberInput, context: WorkflowContext) -> NumberOutput:
        del input_data
        del context.state["flag"]
        return NumberOutput(result=1)

    workflow = Workflow(
        WorkflowConfig(
            name="staged-state-delete",
            state_mode=WorkflowStateMode.STAGED,
            initial_state={"flag": "keep", "other": "value"},
        )
    )
    workflow.set_start_step(
        FunctionStep(
            step_id="delete_flag",
            metadata=StepMetadata(name="delete_flag"),
            input_type=NumberInput,
            output_type=NumberOutput,
            func=delete_flag,
        )
    ).add_end_step("delete_flag")

    execution = await WorkflowRunner().run(workflow, {"value": 1})

    assert execution.status == WorkflowStatus.COMPLETED
    assert "flag" not in execution.state
    assert execution.state["other"] == "value"
