"""Workflow runner with streaming events and bounded parallel execution."""

from __future__ import annotations

import asyncio
import os
import traceback
from contextlib import contextmanager
from dataclasses import dataclass
from datetime import datetime
from importlib import import_module
from typing import Any, AsyncGenerator

from pydantic import TypeAdapter

from agentbyte.cancellation_token import CancellationToken
from agentbyte.workflow.core.checkpoint import (
    CheckpointConfig,
    CheckpointValidationResult,
    WorkflowCheckpoint,
)
from agentbyte.workflow.core.models import (
    CheckpointSavedEvent,
    EdgeActivatedEvent,
    StepCompletedEvent,
    StepExecution,
    StepFailedEvent,
    StepProgressEvent,
    StepStartedEvent,
    StepStatus,
    WorkflowApprovalRequestedEvent,
    WorkflowCancelledEvent,
    WorkflowCompletedEvent,
    WorkflowCompletionMode,
    WorkflowErrorDetails,
    WorkflowEvent,
    WorkflowExecution,
    WorkflowFailedEvent,
    WorkflowInputRequestedEvent,
    WorkflowPendingRequest,
    WorkflowRequestKind,
    WorkflowResumedEvent,
    WorkflowStartedEvent,
    WorkflowStateConflictPolicy,
    WorkflowStateMode,
    WorkflowStatus,
    WorkflowSuspendedEvent,
    WorkflowContext,
    _WorkflowRequestSuspended,
    _STATE_DELETED,
)
from agentbyte.workflow.core.workflow import Workflow
from agentbyte.workflow.steps.step import BaseStep


class WorkflowRunError(RuntimeError):
    """Base exception raised by WorkflowRunner.run()."""


class WorkflowFailedError(WorkflowRunError):
    """Workflow execution failed."""


class WorkflowCancelledError(WorkflowRunError):
    """Workflow execution was cancelled."""


@dataclass
class _StepRunResult:
    output_data: dict[str, Any]
    staged_state: dict[str, Any]
    completed_at: datetime


class WorkflowRunner:
    """Executes workflows with dependency resolution and bounded concurrency."""

    def __init__(self, max_concurrent_steps: int = 5) -> None:
        self.max_concurrent_steps = max_concurrent_steps
        self._execution_semaphore = asyncio.Semaphore(max_concurrent_steps)
        self._cancellation_tokens: dict[str, CancellationToken] = {}
        self._active_workflow_runs: set[int] = set()
        self._suspended_executions: dict[int, WorkflowExecution] = {}
        self._pending_response_types: dict[int, dict[str, Any]] = {}

    async def run(
        self,
        workflow: Workflow,
        initial_input: dict[str, Any] | None = None,
        cancellation_token: CancellationToken | None = None,
        checkpoint: WorkflowCheckpoint | None = None,
        checkpoint_config: CheckpointConfig | None = None,
        *,
        responses: dict[str, Any] | None = None,
    ) -> WorkflowExecution:
        """Run a workflow to completion and return the final execution state."""

        final_execution: WorkflowExecution | None = None
        stream = self.run_stream(
            workflow,
            initial_input,
            cancellation_token=cancellation_token,
            checkpoint=checkpoint,
            checkpoint_config=checkpoint_config,
            responses=responses,
        )
        try:
            async for event in stream:
                if isinstance(event, WorkflowCompletedEvent):
                    final_execution = event.execution
                elif isinstance(event, WorkflowSuspendedEvent):
                    final_execution = event.execution
                elif isinstance(event, WorkflowFailedEvent):
                    final_execution = event.execution
                    raise WorkflowFailedError(event.error)
                elif isinstance(event, WorkflowCancelledEvent):
                    final_execution = event.execution
                    raise WorkflowCancelledError(event.reason)
        finally:
            await stream.aclose()

        if final_execution is None:
            msg = "Workflow finished without a terminal execution state"
            raise RuntimeError(msg)
        return final_execution

    async def run_stream(
        self,
        workflow: Workflow,
        initial_input: dict[str, Any] | None = None,
        cancellation_token: CancellationToken | None = None,
        checkpoint: WorkflowCheckpoint | None = None,
        checkpoint_config: CheckpointConfig | None = None,
        *,
        responses: dict[str, Any] | None = None,
    ) -> AsyncGenerator[WorkflowEvent, None]:
        """Run a workflow and yield structured events in real time."""

        workflow_id = self._get_workflow_id(workflow)
        workflow_key = self._get_workflow_key(workflow)
        payload = initial_input.copy() if initial_input else {}

        if initial_input is not None and responses is not None:
            raise ValueError("Cannot provide both initial_input and responses")
        if workflow_key in self._active_workflow_runs:
            raise RuntimeError(f"Workflow '{workflow_id}' is already running")

        self._active_workflow_runs.add(workflow_key)

        if cancellation_token is not None:
            self._cancellation_tokens[workflow_id] = cancellation_token

        execution: WorkflowExecution | None = None

        with self._start_workflow_span(
            workflow=workflow,
            workflow_id=workflow_id,
            initial_input=payload,
            is_resume=checkpoint is not None or responses is not None,
        ) as workflow_span:
            try:
                if checkpoint is not None:
                    validation = self.validate_checkpoint(workflow, checkpoint)
                    if not validation.can_resume:
                        self._mark_span_error(
                            workflow_span,
                            f"Checkpoint validation failed: {validation.errors}",
                        )
                        yield WorkflowFailedEvent(
                            timestamp=datetime.now(),
                            workflow_id=workflow_id,
                            error_details=self._build_error_details(
                                RuntimeError(
                                    f"Checkpoint validation failed: {validation.errors}"
                                ),
                                workflow_id=workflow_id,
                            ),
                        )
                        return
                    self._restore_step_states(workflow, checkpoint)
                    execution = checkpoint.execution.model_copy(deep=True)
                    self._hydrate_pending_response_types(
                        workflow_key=workflow_key,
                        pending_requests=execution.pending_requests.values(),
                    )
                    execution.status = WorkflowStatus.RUNNING
                    execution.end_time = None
                    yield WorkflowResumedEvent(
                        timestamp=datetime.now(),
                        workflow_id=workflow_id,
                        checkpoint_id=checkpoint.checkpoint_id,
                        completed_steps=list(checkpoint.completed_step_ids),
                        pending_steps=list(checkpoint.pending_step_ids),
                    )
                elif responses is not None:
                    execution = self._suspended_executions.get(workflow_key)
                    if execution is None:
                        raise ValueError(
                            f"No suspended workflow execution available for '{workflow_id}'"
                        )
                    unknown_request_ids = sorted(
                        set(responses) - set(execution.pending_requests)
                    )
                    if unknown_request_ids:
                        raise ValueError(
                            f"Unknown workflow request ids for resume: {unknown_request_ids}"
                        )
                    execution.status = WorkflowStatus.RUNNING
                    execution.end_time = None
                    yield WorkflowResumedEvent(
                        timestamp=datetime.now(),
                        workflow_id=workflow_id,
                        checkpoint_id="in_memory_resume",
                        completed_steps=self._get_completed_step_ids(execution),
                        pending_steps=self._get_pending_step_ids(workflow, execution),
                    )
                else:
                    self._suspended_executions.pop(workflow_key, None)
                    self._pending_response_types.pop(workflow_key, None)
                    yield WorkflowStartedEvent(
                        timestamp=datetime.now(),
                        workflow_id=workflow_id,
                        initial_input=payload,
                    )

                    validation = workflow.validate_workflow()
                    if not validation.is_valid:
                        self._mark_span_error(
                            workflow_span,
                            f"Workflow validation failed: {validation.errors}",
                        )
                        yield WorkflowFailedEvent(
                            timestamp=datetime.now(),
                            workflow_id=workflow_id,
                            error_details=self._build_error_details(
                                RuntimeError(
                                    f"Workflow validation failed: {validation.errors}"
                                ),
                                workflow_id=workflow_id,
                            ),
                        )
                        return

                    if workflow.start_step_id:
                        start_step = workflow.steps[workflow.start_step_id]
                        try:
                            start_step.input_type(**payload)
                        except Exception as exc:
                            self._mark_span_error(
                                workflow_span,
                                (
                                    "Initial input validation failed for start step "
                                    f"'{workflow.start_step_id}': {exc}"
                                ),
                            )
                            yield WorkflowFailedEvent(
                                timestamp=datetime.now(),
                                workflow_id=workflow_id,
                                error_details=self._build_error_details(
                                    exc,
                                    workflow_id=workflow_id,
                                    step_id=workflow.start_step_id,
                                    message=(
                                        "Initial input validation failed for start step "
                                        f"'{workflow.start_step_id}': {exc}"
                                    ),
                                ),
                            )
                            return

                    execution = WorkflowExecution(
                        workflow_id=workflow_id,
                        status=WorkflowStatus.RUNNING,
                        start_time=datetime.now(),
                        state=workflow.config.initial_state.copy(),
                    )
                    execution.state.update(payload)

                async for event in self._execute_workflow_stream(
                    workflow=workflow,
                    execution=execution,
                    initial_input=payload,
                    workflow_key=workflow_key,
                    responses=responses or {},
                    checkpoint_config=checkpoint_config,
                ):
                    yield event

                if execution.status == WorkflowStatus.CANCELLED:
                    self._mark_span_cancelled(workflow_span)
                    return
                if execution.status == WorkflowStatus.SUSPENDED:
                    return

                execution.status = WorkflowStatus.COMPLETED
                execution.end_time = datetime.now()
                self._mark_workflow_span_success(workflow_span, workflow, execution)
                # Resolve the terminal step's output for direct frontend consumption.
                # Prefer end_step_ids when declared; fall back to insertion-order last.
                final_output: dict[str, Any] | None = None
                for step_id in workflow.end_step_ids or []:
                    step_exec = execution.step_executions.get(step_id)
                    if step_exec and step_exec.output_data:
                        final_output = step_exec.output_data
                        break
                if final_output is None and execution.step_executions:
                    final_output = list(execution.step_executions.values())[-1].output_data
                yield WorkflowCompletedEvent(
                    timestamp=datetime.now(),
                    workflow_id=workflow_id,
                    execution=execution,
                    final_output=final_output,
                )
            except Exception as exc:
                if execution is not None:
                    execution.status = WorkflowStatus.FAILED
                    execution.error = str(exc)
                    execution.end_time = datetime.now()
                self._mark_span_error(workflow_span, exc)
                yield WorkflowFailedEvent(
                    timestamp=datetime.now(),
                    workflow_id=workflow_id,
                    error_details=self._build_error_details(
                        exc,
                        workflow_id=workflow_id,
                        execution=execution,
                    ),
                    execution=execution,
                )
            finally:
                if execution is not None and execution.pending_requests:
                    self._suspended_executions[workflow_key] = execution
                else:
                    self._suspended_executions.pop(workflow_key, None)
                    self._pending_response_types.pop(workflow_key, None)
                self._cancellation_tokens.pop(workflow_id, None)
                self._active_workflow_runs.discard(workflow_key)

    async def _execute_workflow_stream(
        self,
        workflow: Workflow,
        execution: WorkflowExecution,
        initial_input: dict[str, Any],
        workflow_key: int,
        responses: dict[str, Any],
        checkpoint_config: CheckpointConfig | None = None,
    ) -> AsyncGenerator[WorkflowEvent, None]:
        completed_steps: set[str] = {
            step_id
            for step_id, step_execution in execution.step_executions.items()
            if step_execution.status in {StepStatus.COMPLETED, StepStatus.FAILED_TOLERATED}
        }
        running_tasks: dict[str, asyncio.Task[_StepRunResult]] = {}
        progress_queue: asyncio.Queue[tuple[str, dict[str, Any]]] = asyncio.Queue()
        workflow_id = self._get_workflow_id(workflow)
        steps_since_last_checkpoint = 0
        staged_mode = workflow.config.state_mode == WorkflowStateMode.STAGED
        wave_completed_results: dict[str, dict[str, Any]] = {}
        wave_staged_state: dict[str, dict[str, Any]] = {}
        wave_completion_order: list[str] = []
        wave_completion_timestamps: dict[str, datetime] = {}

        while True:
            cancellation_token = self._cancellation_tokens.get(workflow_id)

            if cancellation_token and cancellation_token.is_cancelled() and not running_tasks:
                execution.status = WorkflowStatus.CANCELLED
                execution.end_time = datetime.now()
                yield WorkflowCancelledEvent(
                    timestamp=datetime.now(),
                    workflow_id=workflow_id,
                    execution=execution,
                    reason="Cancelled by user",
                )
                return

            if workflow.end_step_ids:
                if workflow.config.completion_mode == WorkflowCompletionMode.ALL_END:
                    if all(step_id in completed_steps for step_id in workflow.end_step_ids):
                        break
                else:  # ANY_END (default)
                    if any(step_id in completed_steps for step_id in workflow.end_step_ids):
                        break
            if len(completed_steps) == len(workflow.steps):
                break

            if (
                not (cancellation_token and cancellation_token.is_cancelled())
                and (not staged_mode or not running_tasks)
            ):
                available_resume_slots = max(
                    0, self.max_concurrent_steps - len(running_tasks)
                )
                for request in list(execution.pending_requests.values()):
                    if available_resume_slots <= 0:
                        break
                    if request.request_id not in responses:
                        continue
                    step_id = request.step_id
                    if step_id in running_tasks:
                        continue
                    step_execution = execution.step_executions.get(step_id)
                    if step_execution is None or step_execution.status != StepStatus.SUSPENDED:
                        continue

                    coerced_response = self._coerce_resume_response(
                        workflow_key=workflow_key,
                        request=request,
                        response_value=responses[request.request_id],
                    )
                    step_execution.status = StepStatus.RUNNING
                    step_execution.start_time = datetime.now()
                    yield StepStartedEvent(
                        timestamp=datetime.now(),
                        workflow_id=workflow_id,
                        step_id=step_id,
                        input_data=step_execution.input_data or {},
                    )
                    task = asyncio.create_task(
                        self._run_step_with_semaphore(
                            workflow=workflow,
                            step=workflow.steps[step_id],
                            input_data=step_execution.input_data or {},
                            workflow_state=execution.state,
                            state_mode=workflow.config.state_mode,
                            progress_queue=progress_queue,
                            pending_request=request,
                            response=coerced_response,
                        )
                    )
                    if cancellation_token is not None:
                        cancellation_token.link_future(task)
                    running_tasks[step_id] = task
                    available_resume_slots -= 1

            ready_steps = [
                step
                for step in workflow.get_ready_steps(execution)
                if step.step_id not in completed_steps and step.step_id not in running_tasks
            ]

            if (
                not (cancellation_token and cancellation_token.is_cancelled())
                and (not staged_mode or not running_tasks)
            ):
                available_slots = max(0, self.max_concurrent_steps - len(running_tasks))
                for step in ready_steps[:available_slots]:
                    step_id = step.step_id
                    input_data = self._prepare_step_input(
                        step_id=step_id,
                        workflow=workflow,
                        execution=execution,
                        initial_input=initial_input,
                    )

                    execution.step_executions[step_id] = StepExecution(
                        step_id=step_id,
                        status=StepStatus.RUNNING,
                        start_time=datetime.now(),
                        input_data=input_data,
                    )
                    yield StepStartedEvent(
                        timestamp=datetime.now(),
                        workflow_id=workflow_id,
                        step_id=step_id,
                        input_data=input_data,
                    )

                    task = asyncio.create_task(
                        self._run_step_with_semaphore(
                            workflow=workflow,
                            step=step,
                            input_data=input_data,
                            workflow_state=execution.state,
                            state_mode=workflow.config.state_mode,
                            progress_queue=progress_queue,
                        )
                    )
                    if cancellation_token is not None:
                        cancellation_token.link_future(task)
                    running_tasks[step_id] = task

            if execution.pending_requests and not running_tasks:
                execution.status = WorkflowStatus.SUSPENDED
                execution.end_time = datetime.now()
                if checkpoint_config is not None and checkpoint_config.auto_save:
                    checkpoint_obj = self._create_checkpoint(
                        workflow=workflow,
                        execution=execution,
                        checkpoint_type="suspend",
                    )
                    await checkpoint_config.store.save(checkpoint_obj)
                    yield CheckpointSavedEvent(
                        timestamp=datetime.now(),
                        workflow_id=workflow_id,
                        checkpoint_id=checkpoint_obj.checkpoint_id,
                        completed_steps=len(completed_steps),
                        total_steps=len(workflow.steps),
                    )
                    if checkpoint_config.auto_cleanup:
                        await checkpoint_config.store.cleanup_old(
                            workflow_id=workflow_id,
                            keep_last_n=checkpoint_config.keep_last_n,
                        )
                yield WorkflowSuspendedEvent(
                    timestamp=datetime.now(),
                    workflow_id=workflow_id,
                    execution=execution,
                    reason="Waiting for external workflow input",
                )
                return

            if not ready_steps and not running_tasks:
                remaining_steps = set(workflow.steps) - completed_steps
                if remaining_steps:
                    msg = (
                        "Workflow stuck: remaining steps "
                        f"{sorted(remaining_steps)} cannot be executed"
                    )
                    raise RuntimeError(msg)
                break

            if not running_tasks:
                continue

            done, _ = await asyncio.wait(
                running_tasks.values(),
                return_when=asyncio.FIRST_COMPLETED,
                timeout=0.01,
            )

            while not progress_queue.empty():
                step_id, progress_data = progress_queue.get_nowait()
                yield StepProgressEvent(
                    timestamp=datetime.now(),
                    workflow_id=workflow_id,
                    step_id=step_id,
                    message=progress_data["message"],
                    completed=progress_data.get("completed"),
                    total=progress_data.get("total"),
                    metadata=progress_data.get("metadata", {}),
                )

            if not done:
                continue

            for task in done:
                step_id = next(
                    sid for sid, running_task in running_tasks.items() if running_task is task
                )
                step_execution = execution.step_executions[step_id]

                try:
                    result = await task
                    step_execution.status = StepStatus.COMPLETED
                    step_execution.output_data = result.output_data
                    step_execution.end_time = datetime.now()
                    if step_execution.pending_request_id:
                        execution.pending_requests.pop(step_execution.pending_request_id, None)
                        self._pending_response_types.get(workflow_key, {}).pop(
                            step_execution.pending_request_id, None
                        )
                        step_execution.pending_request_id = None
                    completed_steps.add(step_id)
                    steps_since_last_checkpoint += 1
                    if staged_mode:
                        wave_completed_results[step_id] = result.output_data
                        wave_staged_state[step_id] = result.staged_state
                        wave_completion_order.append(step_id)
                        wave_completion_timestamps[step_id] = result.completed_at
                    else:
                        execution.state[f"{step_id}_output"] = result.output_data

                    yield StepCompletedEvent(
                        timestamp=datetime.now(),
                        workflow_id=workflow_id,
                        step_id=step_id,
                        output_data=result.output_data,
                        duration_seconds=self._duration_seconds(step_execution),
                    )

                    if not staged_mode:
                        for edge in workflow.edges:
                            if edge.from_step != step_id:
                                continue
                            if workflow._evaluate_edge_condition(edge, execution):
                                yield EdgeActivatedEvent(
                                    timestamp=datetime.now(),
                                    workflow_id=workflow_id,
                                    from_step=step_id,
                                    to_step=edge.to_step,
                                    data=result.output_data,
                                )

                    if (
                        not staged_mode
                        and checkpoint_config is not None
                        and checkpoint_config.auto_save
                        and steps_since_last_checkpoint >= checkpoint_config.save_interval_steps
                    ):
                        checkpoint_obj = self._create_checkpoint(
                            workflow=workflow,
                            execution=execution,
                            checkpoint_type="auto",
                        )
                        await checkpoint_config.store.save(checkpoint_obj)
                        yield CheckpointSavedEvent(
                            timestamp=datetime.now(),
                            workflow_id=workflow_id,
                            checkpoint_id=checkpoint_obj.checkpoint_id,
                            completed_steps=len(completed_steps),
                            total_steps=len(workflow.steps),
                        )
                        if checkpoint_config.auto_cleanup:
                            await checkpoint_config.store.cleanup_old(
                                workflow_id=workflow_id,
                                keep_last_n=checkpoint_config.keep_last_n,
                            )
                        steps_since_last_checkpoint = 0
                except asyncio.CancelledError:
                    step_execution.status = StepStatus.CANCELLED
                    step_execution.end_time = datetime.now()
                    step_execution.error = "Step was cancelled"
                    yield StepFailedEvent(
                        timestamp=datetime.now(),
                        workflow_id=workflow_id,
                        step_id=step_id,
                        error_details=self._build_error_details(
                            asyncio.CancelledError("Step was cancelled"),
                            workflow_id=workflow_id,
                            step_id=step_id,
                            execution=execution,
                        ),
                        duration_seconds=self._duration_seconds(step_execution),
                    )
                except _WorkflowRequestSuspended as request_signal:
                    step_execution.status = StepStatus.SUSPENDED
                    step_execution.end_time = datetime.now()
                    step_execution.pending_request_id = request_signal.request.request_id
                    execution.pending_requests[
                        request_signal.request.request_id
                    ] = request_signal.request
                    self._pending_response_types.setdefault(workflow_key, {})[
                        request_signal.request.request_id
                    ] = request_signal.response_type
                    if request_signal.request.kind == WorkflowRequestKind.INPUT:
                        yield WorkflowInputRequestedEvent(
                            timestamp=datetime.now(),
                            workflow_id=workflow_id,
                            step_id=step_id,
                            request=request_signal.request,
                        )
                    else:
                        yield WorkflowApprovalRequestedEvent(
                            timestamp=datetime.now(),
                            workflow_id=workflow_id,
                            step_id=step_id,
                            request=request_signal.request,
                        )
                except Exception as exc:
                    step_execution.status = StepStatus.FAILED
                    step_execution.end_time = datetime.now()
                    step_execution.error = str(exc)
                    yield StepFailedEvent(
                        timestamp=datetime.now(),
                        workflow_id=workflow_id,
                        step_id=step_id,
                        error_details=self._build_error_details(
                            exc,
                            workflow_id=workflow_id,
                            step_id=step_id,
                            execution=execution,
                        ),
                        duration_seconds=self._duration_seconds(step_execution),
                    )
                    # Tolerate branch failure when ALL downstream join steps are SKIP_FAILED
                    downstream_ids = [
                        edge.to_step for edge in workflow.edges if edge.from_step == step_id
                    ]
                    if downstream_ids and all(
                        workflow.is_step_failure_tolerated(dsid) for dsid in downstream_ids
                    ):
                        step_execution.status = StepStatus.FAILED_TOLERATED
                        completed_steps.add(step_id)
                    else:
                        raise
                finally:
                    del running_tasks[step_id]

            if staged_mode and not running_tasks and wave_completion_order:
                checkpoint_obj = self._commit_staged_wave(
                    workflow=workflow,
                    execution=execution,
                    checkpoint_config=checkpoint_config,
                    steps_since_last_checkpoint=steps_since_last_checkpoint,
                    workflow_id=workflow_id,
                    wave_completion_order=wave_completion_order,
                    wave_completion_timestamps=wave_completion_timestamps,
                    wave_completed_results=wave_completed_results,
                    wave_staged_state=wave_staged_state,
                )
                if checkpoint_obj is not None:
                    await checkpoint_config.store.save(checkpoint_obj)
                    yield CheckpointSavedEvent(
                        timestamp=datetime.now(),
                        workflow_id=workflow_id,
                        checkpoint_id=checkpoint_obj.checkpoint_id,
                        completed_steps=len(completed_steps),
                        total_steps=len(workflow.steps),
                    )
                    if checkpoint_config.auto_cleanup:
                        await checkpoint_config.store.cleanup_old(
                            workflow_id=workflow_id,
                            keep_last_n=checkpoint_config.keep_last_n,
                        )
                steps_since_last_checkpoint = 0
                for step_id in wave_completion_order:
                    result = wave_completed_results[step_id]
                    for edge in workflow.edges:
                        if edge.from_step != step_id:
                            continue
                        if workflow._evaluate_edge_condition(edge, execution):
                            yield EdgeActivatedEvent(
                                timestamp=datetime.now(),
                                workflow_id=workflow_id,
                                from_step=step_id,
                                to_step=edge.to_step,
                                data=result,
                            )
                wave_completed_results.clear()
                wave_staged_state.clear()
                wave_completion_order.clear()
                wave_completion_timestamps.clear()

    async def _run_step_with_semaphore(
        self,
        workflow: Workflow,
        step: BaseStep[Any, Any],
        input_data: dict[str, Any],
        workflow_state: dict[str, Any],
        state_mode: WorkflowStateMode,
        progress_queue: asyncio.Queue[tuple[str, dict[str, Any]]],
        pending_request: WorkflowPendingRequest | None = None,
        response: Any | None = None,
    ) -> _StepRunResult:
        async with self._execution_semaphore:
            staged_state: dict[str, Any] = {}
            if state_mode == WorkflowStateMode.STAGED:
                context = WorkflowContext.from_staged_state(
                    workflow_state,
                    staged_writes=staged_state,
                    step_id=step.step_id,
                )
            else:
                context = WorkflowContext.from_state_ref(workflow_state, step_id=step.step_id)

            def progress_callback(progress_data: dict[str, Any]) -> None:
                progress_queue.put_nowait((step.step_id, progress_data))

            context._progress_callback = progress_callback
            with self._start_step_span(workflow, step, input_data) as step_span:
                try:
                    result = await step.run(
                        input_data,
                        context,
                        pending_request=pending_request,
                        response=response,
                    )
                    self._mark_step_span_success(step_span, result)
                    return _StepRunResult(
                        output_data=result,
                        staged_state=staged_state.copy(),
                        completed_at=datetime.now(),
                    )
                except asyncio.CancelledError:
                    self._mark_span_cancelled(step_span)
                    raise
                except _WorkflowRequestSuspended:
                    raise
                except Exception as exc:
                    self._mark_span_error(step_span, exc, is_step=True)
                    raise

    def _commit_staged_wave(
        self,
        *,
        workflow: Workflow,
        execution: WorkflowExecution,
        checkpoint_config: CheckpointConfig | None,
        steps_since_last_checkpoint: int,
        workflow_id: str,
        wave_completion_order: list[str],
        wave_completion_timestamps: dict[str, datetime],
        wave_completed_results: dict[str, dict[str, Any]],
        wave_staged_state: dict[str, dict[str, Any]],
    ) -> WorkflowCheckpoint | None:
        del workflow_id
        wave_completion_order[:] = sorted(
            wave_completion_order or list(wave_staged_state),
            key=lambda step_id: wave_completion_timestamps[step_id],
        )
        merged_writes = self._merge_staged_writes(
            wave_staged_state,
            conflict_policy=workflow.config.state_conflict_policy,
            completion_order=wave_completion_order,
        )
        for key, value in merged_writes.items():
            if value is _STATE_DELETED:
                execution.state.pop(key, None)
            else:
                execution.state[key] = value
        for step_id in wave_completion_order:
            execution.state[f"{step_id}_output"] = wave_completed_results[step_id]

        if (
            checkpoint_config is not None
            and checkpoint_config.auto_save
            and steps_since_last_checkpoint >= checkpoint_config.save_interval_steps
        ):
            return self._create_checkpoint(
                workflow=workflow,
                execution=execution,
                checkpoint_type="auto",
            )
        return None

    @staticmethod
    def _merge_staged_writes(
        writes_by_step: dict[str, dict[str, Any]],
        *,
        conflict_policy: WorkflowStateConflictPolicy,
        completion_order: list[str] | None = None,
    ) -> dict[str, Any]:
        merged: dict[str, Any] = {}
        owners: dict[str, str] = {}
        ordered_step_ids = completion_order or sorted(writes_by_step)
        for step_id in ordered_step_ids:
            if step_id not in writes_by_step:
                continue
            for key, value in writes_by_step[step_id].items():
                if key not in merged:
                    merged[key] = value
                    owners[key] = step_id
                    continue
                if merged[key] == value:
                    continue
                if conflict_policy == WorkflowStateConflictPolicy.LAST_WRITE_WINS:
                    merged[key] = value
                    owners[key] = step_id
                    continue
                msg = (
                    "Staged workflow state conflict for key "
                    f"'{key}' between steps '{owners[key]}' and '{step_id}'"
                )
                raise RuntimeError(msg)
        return merged

    def _prepare_step_input(
        self,
        step_id: str,
        workflow: Workflow,
        execution: WorkflowExecution,
        initial_input: dict[str, Any],
    ) -> dict[str, Any]:
        if step_id == workflow.start_step_id:
            return initial_input.copy()

        dependencies = workflow.get_step_dependencies(step_id)

        # Identify always-edge dependencies (these form the fan-in set)
        always_dep_ids = [
            dep for dep in dependencies
            if any(
                edge.from_step == dep and edge.to_step == step_id and edge.condition.type == "always"
                for edge in workflow.edges
            )
        ]

        # Fan-in: ≥2 always-edge dependencies → collect all branch outputs into an envelope
        if len(always_dep_ids) >= 2:
            branch_outputs: dict[str, dict[str, Any]] = {}
            failed_branches: list[str] = []
            for dep_id in always_dep_ids:
                dep_exec = execution.step_executions.get(dep_id)
                if (
                    dep_exec is not None
                    and dep_exec.status == StepStatus.COMPLETED
                    and dep_exec.output_data is not None
                ):
                    branch_outputs[dep_id] = dep_exec.output_data.copy()
                else:
                    failed_branches.append(dep_id)
            return {"branch_outputs": branch_outputs, "failed_branches": failed_branches}

        # Single dependency: pass through the upstream output directly (original behaviour)
        for dependency_id in dependencies:
            dependency_execution = execution.step_executions.get(dependency_id)
            if (
                dependency_execution is not None
                and dependency_execution.status == StepStatus.COMPLETED
                and dependency_execution.output_data is not None
            ):
                return dependency_execution.output_data.copy()

        return initial_input.copy()

    def cancel_workflow(self, workflow_id: str) -> bool:
        """Cancel an in-flight workflow execution if one is registered."""

        cancellation_token = self._cancellation_tokens.get(workflow_id)
        if cancellation_token is None:
            return False
        cancellation_token.cancel()
        return True

    @staticmethod
    def _get_workflow_id(workflow: Workflow) -> str:
        return workflow.id

    @staticmethod
    def _get_workflow_key(workflow: Workflow) -> int:
        return id(workflow)

    def _coerce_resume_response(
        self,
        *,
        workflow_key: int,
        request: WorkflowPendingRequest,
        response_value: Any,
    ) -> Any:
        response_type = self._pending_response_types.get(workflow_key, {}).get(request.request_id)
        if response_type is None:
            response_type = self._resolve_pending_request_response_type(request)
            self._pending_response_types.setdefault(workflow_key, {})[request.request_id] = (
                response_type
            )
        if response_type is None:
            msg = (
                f"Missing response type information for workflow request "
                f"'{request.request_id}'"
            )
            raise RuntimeError(msg)
        return TypeAdapter(response_type).validate_python(response_value)

    @staticmethod
    def _get_completed_step_ids(execution: WorkflowExecution) -> list[str]:
        return [
            step_id
            for step_id, step_execution in execution.step_executions.items()
            if step_execution.status in {StepStatus.COMPLETED, StepStatus.FAILED_TOLERATED}
        ]

    @staticmethod
    def _get_pending_step_ids(
        workflow: Workflow,
        execution: WorkflowExecution,
    ) -> list[str]:
        return [
            step_id
            for step_id in workflow.steps
            if step_id not in execution.step_executions
            or execution.step_executions[step_id].status
            in {StepStatus.PENDING, StepStatus.SUSPENDED}
        ]

    @staticmethod
    def _duration_seconds(step_execution: StepExecution) -> float:
        if step_execution.start_time and step_execution.end_time:
            return (step_execution.end_time - step_execution.start_time).total_seconds()
        return 0.0

    def validate_checkpoint(
        self,
        workflow: Workflow,
        checkpoint: WorkflowCheckpoint,
    ) -> CheckpointValidationResult:
        """Validate checkpoint compatibility before resume."""

        errors: list[str] = []
        workflow_id = self._get_workflow_id(workflow)
        if checkpoint.workflow_id != workflow_id:
            errors.append(
                f"Checkpoint workflow_id '{checkpoint.workflow_id}' does not match '{workflow_id}'"
            )
        if checkpoint.workflow_structure_hash != workflow.compute_structure_hash():
            errors.append("Workflow structure changed since checkpoint was created")

        return CheckpointValidationResult(
            is_valid=not errors,
            errors=errors,
            can_resume=not errors,
            checkpoint_info={
                "checkpoint_id": checkpoint.checkpoint_id,
                "completed_steps": len(checkpoint.completed_step_ids),
                "pending_steps": len(checkpoint.pending_step_ids),
                "workflow_version": checkpoint.workflow_version,
            },
        )

    def _create_checkpoint(
        self,
        workflow: Workflow,
        execution: WorkflowExecution,
        checkpoint_type: str = "manual",
    ) -> WorkflowCheckpoint:
        version = str(workflow.config.metadata.get("version", "1.0.0"))
        return WorkflowCheckpoint.from_execution(
            execution=execution,
            workflow_id=self._get_workflow_id(workflow),
            workflow_version=version,
            workflow_structure_hash=workflow.compute_structure_hash(),
            all_step_ids=list(workflow.steps.keys()),
            checkpoint_type=checkpoint_type,
            step_states=self._collect_step_states(workflow),
        )

    @staticmethod
    def _collect_step_states(workflow: Workflow) -> dict[str, dict[str, Any]]:
        step_states: dict[str, dict[str, Any]] = {}
        for step_id, step in workflow.steps.items():
            state = step.on_checkpoint_save()
            if state:
                step_states[step_id] = state.copy()
        return step_states

    @staticmethod
    def _restore_step_states(
        workflow: Workflow,
        checkpoint: WorkflowCheckpoint,
    ) -> None:
        for step_id, state in checkpoint.step_states.items():
            step = workflow.steps.get(step_id)
            if step is None:
                continue
            step.on_checkpoint_restore(state.copy())

    def _hydrate_pending_response_types(
        self,
        *,
        workflow_key: int,
        pending_requests: Any,
    ) -> None:
        resolved_types: dict[str, Any] = {}
        for request in pending_requests:
            try:
                resolved_types[request.request_id] = (
                    self._resolve_pending_request_response_type(request)
                )
            except RuntimeError:
                continue
        if resolved_types:
            self._pending_response_types[workflow_key] = resolved_types

    @staticmethod
    def _resolve_pending_request_response_type(request: WorkflowPendingRequest) -> Any:
        module_name = request.response_type_module
        qualname = request.response_type_qualname
        if not module_name or not qualname:
            return None
        try:
            target: Any = import_module(module_name)
            for attribute in qualname.split("."):
                if attribute == "<locals>":
                    msg = (
                        "Cannot restore a locally defined response type "
                        f"for workflow request '{request.request_id}'"
                    )
                    raise RuntimeError(msg)
                target = getattr(target, attribute)
            return target
        except RuntimeError:
            raise
        except Exception as exc:
            msg = (
                "Unable to restore response type "
                f"'{request.response_type_name}' for workflow request "
                f"'{request.request_id}' from '{module_name}.{qualname}'"
            )
            raise RuntimeError(msg) from exc

    @staticmethod
    def _build_error_details(
        error: BaseException,
        *,
        workflow_id: str,
        step_id: str | None = None,
        execution: WorkflowExecution | None = None,
        message: str | None = None,
    ) -> WorkflowErrorDetails:
        formatted_traceback = None
        if error.__traceback__ is not None:
            formatted_traceback = "".join(
                traceback.format_exception(type(error), error, error.__traceback__)
            )
        return WorkflowErrorDetails(
            error_type=type(error).__name__,
            message=message or str(error),
            traceback=formatted_traceback,
            workflow_id=workflow_id,
            step_id=step_id,
            execution_id=execution.id if execution is not None else None,
        )

    @staticmethod
    def _otel_enabled() -> bool:
        return os.getenv("AGENTBYTE_ENABLE_OTEL", "false").lower() in ("true", "1", "yes")

    @contextmanager
    def _start_workflow_span(
        self,
        workflow: Workflow,
        workflow_id: str,
        initial_input: dict[str, Any],
        is_resume: bool,
    ) -> Any:
        if not self._otel_enabled():
            yield None
            return

        try:
            from opentelemetry import trace

            tracer = trace.get_tracer("agentbyte.workflow")
            with tracer.start_as_current_span(f"workflow {workflow.config.name}") as span:
                span.set_attribute("agentbyte.workflow.id", workflow_id)
                span.set_attribute("agentbyte.workflow.name", workflow.config.name)
                span.set_attribute("agentbyte.workflow.step_count", len(workflow.steps))
                span.set_attribute("agentbyte.workflow.edge_count", len(workflow.edges))
                span.set_attribute("agentbyte.workflow.end_step_count", len(workflow.end_step_ids))
                span.set_attribute("agentbyte.workflow.resumed", is_resume)
                span.set_attribute("agentbyte.workflow.run_mode", "stream")
                span.set_attribute(
                    "agentbyte.workflow.input_keys",
                    ",".join(sorted(initial_input.keys())),
                )
                yield span
        except Exception:
            yield None

    @contextmanager
    def _start_step_span(
        self,
        workflow: Workflow,
        step: BaseStep[Any, Any],
        input_data: dict[str, Any],
    ) -> Any:
        if not self._otel_enabled():
            yield None
            return

        try:
            from opentelemetry import trace

            tracer = trace.get_tracer("agentbyte.workflow")
            with tracer.start_as_current_span(f"workflow step {step.step_id}") as span:
                span.set_attribute("agentbyte.workflow.id", self._get_workflow_id(workflow))
                span.set_attribute("agentbyte.workflow.name", workflow.config.name)
                span.set_attribute("agentbyte.workflow.step.id", step.step_id)
                span.set_attribute(
                    "agentbyte.workflow.step.name",
                    step.metadata.name,
                )
                span.set_attribute(
                    "agentbyte.workflow.step.type",
                    step.__class__.__name__,
                )
                span.set_attribute(
                    "agentbyte.workflow.step.input_keys",
                    ",".join(sorted(input_data.keys())),
                )
                yield span
        except Exception:
            yield None

    @staticmethod
    def _mark_span_error(
        span: Any,
        error: Any,
        *,
        is_step: bool = False,
    ) -> None:
        if span is None:
            return

        status_attr = (
            "agentbyte.workflow.step.status"
            if is_step
            else "agentbyte.workflow.status"
        )
        span.set_attribute(status_attr, "failed")
        span.set_attribute("agentbyte.workflow.error", str(error))

        try:
            from opentelemetry.trace import Status, StatusCode

            if isinstance(error, Exception):
                span.record_exception(error)
            span.set_status(Status(StatusCode.ERROR, str(error)))
        except Exception:
            return

    @staticmethod
    def _mark_span_cancelled(span: Any) -> None:
        if span is None:
            return

        span.set_attribute("agentbyte.workflow.status", "cancelled")

    @staticmethod
    def _mark_workflow_span_success(
        span: Any,
        workflow: Workflow,
        execution: WorkflowExecution,
    ) -> None:
        if span is None:
            return

        span.set_attribute("agentbyte.workflow.status", execution.status.value)
        span.set_attribute(
            "agentbyte.workflow.completed_steps",
            len(
                [
                    step_execution
                    for step_execution in execution.step_executions.values()
                    if step_execution.status == StepStatus.COMPLETED
                ]
            ),
        )
        span.set_attribute("agentbyte.workflow.total_steps", len(workflow.steps))

        try:
            from opentelemetry.trace import Status, StatusCode

            span.set_status(Status(StatusCode.OK))
        except Exception:
            return

    @staticmethod
    def _mark_step_span_success(
        span: Any,
        result: dict[str, Any],
    ) -> None:
        if span is None:
            return

        span.set_attribute("agentbyte.workflow.step.status", "completed")
        span.set_attribute(
            "agentbyte.workflow.step.output_keys",
            ",".join(sorted(result.keys())),
        )

        try:
            from opentelemetry.trace import Status, StatusCode

            span.set_status(Status(StatusCode.OK))
        except Exception:
            return


__all__ = ["WorkflowRunner"]
