import asyncio
import json
from collections.abc import Callable
from datetime import datetime, timedelta
from functools import cached_property
from typing import (
    TYPE_CHECKING,
    Any,
    Concatenate,
    Generic,
    Literal,
    ParamSpec,
    TypeVar,
    cast,
    get_type_hints,
    overload,
)

from google.protobuf import timestamp_pb2
from pydantic import BaseModel, ConfigDict, SkipValidation, TypeAdapter, model_validator

from hatchet_sdk.clients.admin import (
    ScheduleTriggerWorkflowOptions,
    TriggerWorkflowOptions,
    WorkflowRunTriggerConfig,
)
from hatchet_sdk.clients.listeners.run_event_listener import RunEventListener
from hatchet_sdk.clients.rest.models.cron_workflows import CronWorkflows
from hatchet_sdk.clients.rest.models.v1_filter import V1Filter
from hatchet_sdk.clients.rest.models.v1_task_status import V1TaskStatus
from hatchet_sdk.clients.rest.models.v1_task_summary import V1TaskSummary
from hatchet_sdk.conditions import Condition, OrGroup
from hatchet_sdk.context.context import Context, DurableContext
from hatchet_sdk.contracts.v1.workflows_pb2 import (
    CreateWorkflowVersionRequest,
    DesiredWorkerLabels,
)
from hatchet_sdk.contracts.v1.workflows_pb2 import StickyStrategy as StickyStrategyProto
from hatchet_sdk.contracts.workflows_pb2 import WorkflowVersion
from hatchet_sdk.labels import DesiredWorkerLabel
from hatchet_sdk.rate_limit import RateLimit
from hatchet_sdk.runnables.task import Task
from hatchet_sdk.runnables.types import (
    ConcurrencyExpression,
    EmptyModel,
    R,
    StepType,
    TaskDefaults,
    TaskPayloadForInternalUse,
    TWorkflowInput,
    WorkflowConfig,
    normalize_validator,
)
from hatchet_sdk.serde import HATCHET_PYDANTIC_SENTINEL
from hatchet_sdk.utils.proto_enums import convert_python_enum_to_proto
from hatchet_sdk.utils.timedelta_to_expression import Duration
from hatchet_sdk.utils.typing import CoroutineLike, JSONSerializableMapping
from hatchet_sdk.workflow_run import WorkflowRunRef

if TYPE_CHECKING:
    from hatchet_sdk import Hatchet


T = TypeVar("T")
P = ParamSpec("P")


def fall_back_to_default(value: T, param_default: T, fallback_value: T | None) -> T:
    ## If the value is not the param default, it's set
    if value != param_default:
        return value

    ## Otherwise, it's unset, so return the fallback value if it's set
    if fallback_value is not None:
        return fallback_value

    ## Otherwise return the param value
    return value


class ComputedTaskParameters(BaseModel):
    schedule_timeout: Duration
    execution_timeout: Duration
    retries: int
    backoff_factor: float | None
    backoff_max_seconds: int | None

    task_defaults: TaskDefaults

    @model_validator(mode="after")
    def validate_params(self) -> "ComputedTaskParameters":
        self.execution_timeout = fall_back_to_default(
            value=self.execution_timeout,
            param_default=timedelta(seconds=60),
            fallback_value=self.task_defaults.execution_timeout,
        )
        self.schedule_timeout = fall_back_to_default(
            value=self.schedule_timeout,
            param_default=timedelta(minutes=5),
            fallback_value=self.task_defaults.schedule_timeout,
        )
        self.backoff_factor = fall_back_to_default(
            value=self.backoff_factor,
            param_default=None,
            fallback_value=self.task_defaults.backoff_factor,
        )
        self.backoff_max_seconds = fall_back_to_default(
            value=self.backoff_max_seconds,
            param_default=None,
            fallback_value=self.task_defaults.backoff_max_seconds,
        )
        self.retries = fall_back_to_default(
            value=self.retries,
            param_default=0,
            fallback_value=self.task_defaults.retries,
        )

        return self


def transform_desired_worker_label(d: DesiredWorkerLabel) -> DesiredWorkerLabels:
    value = d.value
    return DesiredWorkerLabels(
        strValue=value if not isinstance(value, int) else None,
        intValue=value if isinstance(value, int) else None,
        required=d.required,
        weight=d.weight,
        comparator=d.comparator,  # type: ignore[arg-type]
    )


class TypedTriggerWorkflowRunConfig(BaseModel, Generic[TWorkflowInput]):
    model_config = ConfigDict(arbitrary_types_allowed=True)
    input: SkipValidation[TWorkflowInput]
    options: TriggerWorkflowOptions


class BaseWorkflow(Generic[TWorkflowInput]):
    def __init__(self, config: WorkflowConfig, client: "Hatchet") -> None:
        self.config = config
        self._default_tasks: list[Task[TWorkflowInput, Any]] = []
        self._durable_tasks: list[Task[TWorkflowInput, Any]] = []
        self._on_failure_task: Task[TWorkflowInput, Any] | None = None
        self._on_success_task: Task[TWorkflowInput, Any] | None = None
        self.client = client

    @property
    def service_name(self) -> str:
        return self.client.config.apply_namespace(self.config.name.lower())

    def _create_action_name(self, step: Task[TWorkflowInput, Any]) -> str:
        return self.service_name + ":" + step.name

    def _is_leaf_task(self, task: Task[TWorkflowInput, Any]) -> bool:
        return not any(task in t.parents for t in self.tasks if task != t)

    def to_proto(self) -> CreateWorkflowVersionRequest:
        namespace = self.client.config.namespace
        service_name = self.service_name

        name = self.name
        event_triggers = [
            self.client.config.apply_namespace(event, namespace)
            for event in self.config.on_events
        ]

        if self._on_success_task:
            self._on_success_task.parents = [
                task
                for task in self.tasks
                if task.type == StepType.DEFAULT and self._is_leaf_task(task)
            ]

        on_success_task = (
            t.to_proto(service_name) if (t := self._on_success_task) else None
        )

        tasks = [
            task.to_proto(service_name)
            for task in self.tasks
            if task.type == StepType.DEFAULT
        ]

        if on_success_task:
            tasks += [on_success_task]

        on_failure_task = (
            t.to_proto(service_name) if (t := self._on_failure_task) else None
        )

        if isinstance(self.config.concurrency, list):
            _concurrency_arr = [c.to_proto() for c in self.config.concurrency]
            _concurrency = None
        elif isinstance(self.config.concurrency, ConcurrencyExpression):
            _concurrency_arr = []
            _concurrency = self.config.concurrency.to_proto()
        elif isinstance(self.config.concurrency, int):
            _concurrency_arr = []
            _concurrency = ConcurrencyExpression.from_int(
                self.config.concurrency
            ).to_proto()
        else:
            _concurrency = None
            _concurrency_arr = []

        # Hack to not send a JSON schema if the input type is None/EmptyModel
        input_type = self.config.input_validator.core_schema.get("cls")

        if input_type is None or input_type is EmptyModel:
            json_schema = None
        else:
            try:
                json_schema = json.dumps(
                    self.config.input_validator.json_schema()
                ).encode("utf-8")
            except Exception:
                json_schema = None

        return CreateWorkflowVersionRequest(
            name=name,
            description=self.config.description,
            version=self.config.version,
            event_triggers=event_triggers,
            cron_triggers=self.config.on_crons,
            tasks=tasks,
            ## TODO: Fix this
            cron_input=None,
            on_failure_task=on_failure_task,
            sticky=convert_python_enum_to_proto(
                self.config.sticky, StickyStrategyProto
            ),  # type: ignore[arg-type]
            concurrency=_concurrency,
            concurrency_arr=_concurrency_arr,
            default_priority=self.config.default_priority,
            default_filters=[f.to_proto() for f in self.config.default_filters],
            input_json_schema=json_schema,
        )

    def _get_workflow_input(self, ctx: Context) -> TWorkflowInput:
        return cast(
            TWorkflowInput,
            self.config.input_validator.validate_python(
                ctx.workflow_input, context=HATCHET_PYDANTIC_SENTINEL
            ),
        )

    def _combine_additional_metadata(
        self, additional_metadata_from_trigger: JSONSerializableMapping
    ) -> JSONSerializableMapping:
        return {
            **self.config.default_additional_metadata,
            **additional_metadata_from_trigger,
        }

    def _create_options_with_combined_additional_meta(
        self, options: TriggerWorkflowOptions
    ) -> TriggerWorkflowOptions:
        options_copy = options.model_copy()
        options_copy.additional_metadata = self._combine_additional_metadata(
            options.additional_metadata
        )

        return options_copy

    @property
    def input_validator(self) -> type[TWorkflowInput]:
        return cast(type[TWorkflowInput], self.config.input_validator)

    @property
    def tasks(self) -> list[Task[TWorkflowInput, Any]]:
        tasks = self._default_tasks + self._durable_tasks

        if self._on_failure_task:
            tasks += [self._on_failure_task]

        if self._on_success_task:
            tasks += [self._on_success_task]

        return tasks

    @property
    def name(self) -> str:
        """
        The (namespaced) name of the workflow.
        """
        return self.client.config.namespace + self.config.name

    def create_bulk_run_item(
        self,
        input: TWorkflowInput = cast(TWorkflowInput, EmptyModel()),
        key: str | None = None,
        options: TriggerWorkflowOptions = TriggerWorkflowOptions(),
    ) -> WorkflowRunTriggerConfig:
        """
        Create a bulk run item for the workflow. This is intended to be used in conjunction with the various `run_many` methods.

        :param input: The input data for the workflow.
        :param key: The key for the workflow run. This is used to identify the run in the bulk operation and for deduplication.
        :param options: Additional options for the workflow run.

        :returns: A `WorkflowRunTriggerConfig` object that can be used to trigger the workflow run, which you then pass into the `run_many` methods.
        """
        return WorkflowRunTriggerConfig(
            workflow_name=self.config.name,
            input=self._serialize_input(input),
            options=self._create_options_with_combined_additional_meta(options),
            key=key,
        )

    def _serialize_input(self, input: TWorkflowInput | None) -> JSONSerializableMapping:
        if not input:
            return {}

        return cast(
            JSONSerializableMapping,
            self.config.input_validator.dump_python(
                input,  # type: ignore[arg-type]
                mode="json",
                context=HATCHET_PYDANTIC_SENTINEL,
            ),
        )

    @cached_property
    def id(self) -> str:
        """
        Get the ID of the workflow.

        :raises ValueError: If no workflow ID is found for the workflow name.
        :returns: The ID of the workflow.
        """
        workflows = self.client.workflows.list(workflow_name=self.name)

        if not workflows.rows:
            raise ValueError(f"No id found for {self.name}")

        for workflow in workflows.rows:
            if workflow.name == self.name:
                return workflow.metadata.id

        raise ValueError(f"No id found for {self.name}")

    def list_runs(
        self,
        since: datetime | None = None,
        until: datetime | None = None,
        limit: int = 100,
        offset: int | None = None,
        statuses: list[V1TaskStatus] | None = None,
        additional_metadata: dict[str, str] | None = None,
        worker_id: str | None = None,
        parent_task_external_id: str | None = None,
        only_tasks: bool = False,
        triggering_event_external_id: str | None = None,
    ) -> list[V1TaskSummary]:
        """
        List runs of the workflow.

        :param since: The start time for the runs to be listed.
        :param until: The end time for the runs to be listed.
        :param limit: The maximum number of runs to be listed.
        :param offset: The offset for pagination.
        :param statuses: The statuses of the runs to be listed.
        :param additional_metadata: Additional metadata for filtering the runs.
        :param worker_id: The ID of the worker that ran the tasks.
        :param parent_task_external_id: The external ID of the parent task.
        :param only_tasks: Whether to list only task runs.
        :param triggering_event_external_id: The event id that triggered the task run.

        :returns: A list of `V1TaskSummary` objects representing the runs of the workflow.
        """
        return self.client.runs.list_with_pagination(
            workflow_ids=[self.id],
            since=since,
            only_tasks=only_tasks,
            offset=offset,
            limit=limit,
            statuses=statuses,
            until=until,
            additional_metadata=additional_metadata,
            worker_id=worker_id,
            parent_task_external_id=parent_task_external_id,
            triggering_event_external_id=triggering_event_external_id,
        )

    async def aio_list_runs(
        self,
        since: datetime | None = None,
        until: datetime | None = None,
        limit: int = 100,
        offset: int | None = None,
        statuses: list[V1TaskStatus] | None = None,
        additional_metadata: dict[str, str] | None = None,
        worker_id: str | None = None,
        parent_task_external_id: str | None = None,
        only_tasks: bool = False,
        triggering_event_external_id: str | None = None,
    ) -> list[V1TaskSummary]:
        """
        List runs of the workflow.

        :param since: The start time for the runs to be listed.
        :param until: The end time for the runs to be listed.
        :param limit: The maximum number of runs to be listed.
        :param offset: The offset for pagination.
        :param statuses: The statuses of the runs to be listed.
        :param additional_metadata: Additional metadata for filtering the runs.
        :param worker_id: The ID of the worker that ran the tasks.
        :param parent_task_external_id: The external ID of the parent task.
        :param only_tasks: Whether to list only task runs.
        :param triggering_event_external_id: The event id that triggered the task run.

        :returns: A list of `V1TaskSummary` objects representing the runs of the workflow.
        """
        return await self.client.runs.aio_list_with_pagination(
            workflow_ids=[self.id],
            since=since,
            only_tasks=only_tasks,
            offset=offset,
            limit=limit,
            statuses=statuses,
            until=until,
            additional_metadata=additional_metadata,
            worker_id=worker_id,
            parent_task_external_id=parent_task_external_id,
            triggering_event_external_id=triggering_event_external_id,
        )

    def create_filter(
        self,
        expression: str,
        scope: str,
        payload: JSONSerializableMapping | None = None,
    ) -> V1Filter:
        """
        Create a new filter.

        :param expression: The expression to evaluate for the filter.
        :param scope: The scope for the filter.
        :param payload: The payload to send with the filter.

        :return: The created filter.
        """
        return self.client.filters.create(
            workflow_id=self.id,
            expression=expression,
            scope=scope,
            payload=payload,
        )

    async def aio_create_filter(
        self,
        expression: str,
        scope: str,
        payload: JSONSerializableMapping | None = None,
    ) -> V1Filter:
        """
        Create a new filter.

        :param expression: The expression to evaluate for the filter.
        :param scope: The scope for the filter.
        :param payload: The payload to send with the filter.

        :return: The created filter.
        """
        return await self.client.filters.aio_create(
            workflow_id=self.id,
            expression=expression,
            scope=scope,
            payload=payload,
        )

    def schedule(
        self,
        run_at: datetime,
        input: TWorkflowInput = cast(TWorkflowInput, EmptyModel()),
        options: ScheduleTriggerWorkflowOptions = ScheduleTriggerWorkflowOptions(),
    ) -> WorkflowVersion:
        """
        Schedule a workflow to run at a specific time.

        :param run_at: The time at which to schedule the workflow.
        :param input: The input data for the workflow.
        :param options: Additional options for workflow execution.
        :returns: A `WorkflowVersion` object representing the scheduled workflow.
        """
        return self.client._client.admin.schedule_workflow(
            name=self.config.name,
            schedules=cast(list[datetime | timestamp_pb2.Timestamp], [run_at]),
            input=self._serialize_input(input),
            options=options,
        )

    async def aio_schedule(
        self,
        run_at: datetime,
        input: TWorkflowInput = cast(TWorkflowInput, EmptyModel()),
        options: ScheduleTriggerWorkflowOptions = ScheduleTriggerWorkflowOptions(),
    ) -> WorkflowVersion:
        """
        Schedule a workflow to run at a specific time.

        :param run_at: The time at which to schedule the workflow.
        :param input: The input data for the workflow.
        :param options: Additional options for workflow execution.
        :returns: A `WorkflowVersion` object representing the scheduled workflow.
        """
        return await self.client._client.admin.aio_schedule_workflow(
            name=self.config.name,
            schedules=cast(list[datetime | timestamp_pb2.Timestamp], [run_at]),
            input=self._serialize_input(input),
            options=options,
        )

    def create_cron(
        self,
        cron_name: str,
        expression: str,
        input: TWorkflowInput = cast(TWorkflowInput, EmptyModel()),
        additional_metadata: JSONSerializableMapping | None = None,
        priority: int | None = None,
    ) -> CronWorkflows:
        """
        Create a cron job for the workflow.

        :param cron_name: The name of the cron job.
        :param expression: The cron expression that defines the schedule for the cron job.
        :param input: The input data for the workflow.
        :param additional_metadata: Additional metadata for the cron job.
        :param priority: The priority of the cron job. Must be between 1 and 3, inclusive.

        :returns: A `CronWorkflows` object representing the created cron job.
        """
        return self.client.cron.create(
            workflow_name=self.config.name,
            cron_name=cron_name,
            expression=expression,
            input=self._serialize_input(input),
            additional_metadata=additional_metadata or {},
            priority=priority,
        )

    async def aio_create_cron(
        self,
        cron_name: str,
        expression: str,
        input: TWorkflowInput = cast(TWorkflowInput, EmptyModel()),
        additional_metadata: JSONSerializableMapping | None = None,
        priority: int | None = None,
    ) -> CronWorkflows:
        """
        Create a cron job for the workflow.

        :param cron_name: The name of the cron job.
        :param expression: The cron expression that defines the schedule for the cron job.
        :param input: The input data for the workflow.
        :param additional_metadata: Additional metadata for the cron job.
        :param priority: The priority of the cron job. Must be between 1 and 3, inclusive.

        :returns: A `CronWorkflows` object representing the created cron job.
        """
        return await self.client.cron.aio_create(
            workflow_name=self.config.name,
            cron_name=cron_name,
            expression=expression,
            input=self._serialize_input(input),
            additional_metadata=additional_metadata or {},
            priority=priority,
        )

    def delete(self) -> None:
        """
        Permanently delete the workflow.

        **DANGEROUS: This will delete a workflow and all of its data**
        """
        self.client.workflows.delete(self.id)

    async def aio_delete(self) -> None:
        """
        Permanently delete the workflow.

        **DANGEROUS: This will delete a workflow and all of its data**
        """
        await self.client.workflows.aio_delete(self.id)


class Workflow(BaseWorkflow[TWorkflowInput]):
    """
    A Hatchet workflow, which allows you to define tasks to be run and perform actions on the workflow.

    Workflows in Hatchet represent coordinated units of work that can be triggered, scheduled, or run on a cron schedule.
    Each workflow can contain multiple tasks that can be arranged in dependencies (DAGs), have customized retry behavior,
    timeouts, concurrency controls, and more.

    Example:
    ```python
    from pydantic import BaseModel
    from hatchet_sdk import Hatchet

    class MyInput(BaseModel):
        name: str

    hatchet = Hatchet()
    workflow = hatchet.workflow("my-workflow", input_type=MyInput)

    @workflow.task()
    def greet(input, ctx):
        return f"Hello, {input.name}!"

    # Run the workflow
    result = workflow.run(MyInput(name="World"))
    ```

    Workflows support various execution patterns including:
    - One-time execution with `run()` or `aio_run()`
    - Scheduled execution with `schedule()`
    - Cron-based recurring execution with `create_cron()`
    - Bulk operations with `run_many()`

    Tasks within workflows can be defined with `@workflow.task()` or `@workflow.durable_task()` decorators
    and can be arranged into complex dependency patterns.
    """

    def run_no_wait(
        self,
        input: TWorkflowInput = cast(TWorkflowInput, EmptyModel()),
        options: TriggerWorkflowOptions = TriggerWorkflowOptions(),
    ) -> WorkflowRunRef:
        """
        Synchronously trigger a workflow run without waiting for it to complete.
        This method is useful for starting a workflow run and immediately returning a reference to the run without blocking while the workflow runs.

        :param input: The input data for the workflow.
        :param options: Additional options for workflow execution.

        :returns: A `WorkflowRunRef` object representing the reference to the workflow run.
        """
        return self.client._client.admin.run_workflow(
            workflow_name=self.config.name,
            input=self._serialize_input(input),
            options=self._create_options_with_combined_additional_meta(options),
        )

    def run(
        self,
        input: TWorkflowInput = cast(TWorkflowInput, EmptyModel()),
        options: TriggerWorkflowOptions = TriggerWorkflowOptions(),
    ) -> dict[str, Any]:
        """
        Run the workflow synchronously and wait for it to complete.

        This method triggers a workflow run, blocks until completion, and returns the final result.

        :param input: The input data for the workflow, must match the workflow's input type.
        :param options: Additional options for workflow execution like metadata and parent workflow ID.

        :returns: The result of the workflow execution as a dictionary.
        """

        ref = self.client._client.admin.run_workflow(
            workflow_name=self.config.name,
            input=self._serialize_input(input),
            options=self._create_options_with_combined_additional_meta(options),
        )

        return ref.result()

    async def aio_run_no_wait(
        self,
        input: TWorkflowInput = cast(TWorkflowInput, EmptyModel()),
        options: TriggerWorkflowOptions = TriggerWorkflowOptions(),
    ) -> WorkflowRunRef:
        """
        Asynchronously trigger a workflow run without waiting for it to complete.
        This method is useful for starting a workflow run and immediately returning a reference to the run without blocking while the workflow runs.

        :param input: The input data for the workflow.
        :param options: Additional options for workflow execution.

        :returns: A `WorkflowRunRef` object representing the reference to the workflow run.
        """

        return await self.client._client.admin.aio_run_workflow(
            workflow_name=self.config.name,
            input=self._serialize_input(input),
            options=self._create_options_with_combined_additional_meta(options),
        )

    async def aio_run(
        self,
        input: TWorkflowInput = cast(TWorkflowInput, EmptyModel()),
        options: TriggerWorkflowOptions = TriggerWorkflowOptions(),
    ) -> dict[str, Any]:
        """
        Run the workflow asynchronously and wait for it to complete.

        This method triggers a workflow run, awaits until completion, and returns the final result.

        :param input: The input data for the workflow, must match the workflow's input type.
        :param options: Additional options for workflow execution like metadata and parent workflow ID.

        :returns: The result of the workflow execution as a dictionary.
        """
        ref = await self.client._client.admin.aio_run_workflow(
            workflow_name=self.config.name,
            input=self._serialize_input(input),
            options=self._create_options_with_combined_additional_meta(options),
        )

        return await ref.aio_result()

    def _get_result(
        self, ref: WorkflowRunRef, return_exceptions: bool
    ) -> dict[str, Any] | BaseException:
        try:
            return ref.result()
        except Exception as e:
            if return_exceptions:
                return e
            raise e

    @overload
    def run_many(
        self,
        workflows: list[WorkflowRunTriggerConfig],
        return_exceptions: Literal[True],
    ) -> list[dict[str, Any] | BaseException]: ...

    @overload
    def run_many(
        self,
        workflows: list[WorkflowRunTriggerConfig],
        return_exceptions: Literal[False] = False,
    ) -> list[dict[str, Any]]: ...

    def run_many(
        self,
        workflows: list[WorkflowRunTriggerConfig],
        return_exceptions: bool = False,
    ) -> list[dict[str, Any]] | list[dict[str, Any] | BaseException]:
        """
        Run a workflow in bulk and wait for all runs to complete.
        This method triggers multiple workflow runs, blocks until all of them complete, and returns the final results.

        :param workflows: A list of `WorkflowRunTriggerConfig` objects, each representing a workflow run to be triggered.
        :param return_exceptions: If `True`, exceptions will be returned as part of the results instead of raising them.
        :returns: A list of results for each workflow run.
        """
        refs = self.client._client.admin.run_workflows(
            workflows=workflows,
        )

        return [self._get_result(ref, return_exceptions) for ref in refs]

    @overload
    async def aio_run_many(
        self,
        workflows: list[WorkflowRunTriggerConfig],
        return_exceptions: Literal[True],
    ) -> list[dict[str, Any] | BaseException]: ...

    @overload
    async def aio_run_many(
        self,
        workflows: list[WorkflowRunTriggerConfig],
        return_exceptions: Literal[False] = False,
    ) -> list[dict[str, Any]]: ...

    async def aio_run_many(
        self,
        workflows: list[WorkflowRunTriggerConfig],
        return_exceptions: bool = False,
    ) -> list[dict[str, Any]] | list[dict[str, Any] | BaseException]:
        """
        Run a workflow in bulk and wait for all runs to complete.
        This method triggers multiple workflow runs, blocks until all of them complete, and returns the final results.

        :param workflows: A list of `WorkflowRunTriggerConfig` objects, each representing a workflow run to be triggered.
        :param return_exceptions: If `True`, exceptions will be returned as part of the results instead of raising them.
        :returns: A list of results for each workflow run.
        """
        refs = await self.client._client.admin.aio_run_workflows(
            workflows=workflows,
        )

        return await asyncio.gather(
            *[ref.aio_result() for ref in refs], return_exceptions=return_exceptions
        )

    def run_many_no_wait(
        self,
        workflows: list[WorkflowRunTriggerConfig],
    ) -> list[WorkflowRunRef]:
        """
        Run a workflow in bulk without waiting for all runs to complete.

        This method triggers multiple workflow runs and immediately returns a list of references to the runs without blocking while the workflows run.

        :param workflows: A list of `WorkflowRunTriggerConfig` objects, each representing a workflow run to be triggered.
        :returns: A list of `WorkflowRunRef` objects, each representing a reference to a workflow run.
        """
        return self.client._client.admin.run_workflows(
            workflows=workflows,
        )

    async def aio_run_many_no_wait(
        self,
        workflows: list[WorkflowRunTriggerConfig],
    ) -> list[WorkflowRunRef]:
        """
        Run a workflow in bulk without waiting for all runs to complete.

        This method triggers multiple workflow runs and immediately returns a list of references to the runs without blocking while the workflows run.

        :param workflows: A list of `WorkflowRunTriggerConfig` objects, each representing a workflow run to be triggered.

        :returns: A list of `WorkflowRunRef` objects, each representing a reference to a workflow run.
        """
        return await self.client._client.admin.aio_run_workflows(
            workflows=workflows,
        )

    def _parse_task_name(
        self,
        name: str | None,
        func: Callable[..., Any],
    ) -> str:
        non_null_name = name or func.__name__
        return non_null_name.lower()

    def task(
        self,
        name: str | None = None,
        schedule_timeout: Duration = timedelta(minutes=5),
        execution_timeout: Duration = timedelta(seconds=60),
        parents: list[Task[TWorkflowInput, Any]] | None = None,
        retries: int = 0,
        rate_limits: list[RateLimit] | None = None,
        desired_worker_labels: dict[str, DesiredWorkerLabel] | None = None,
        backoff_factor: float | None = None,
        backoff_max_seconds: int | None = None,
        concurrency: int | list[ConcurrencyExpression] | None = None,
        wait_for: list[Condition | OrGroup] | None = None,
        skip_if: list[Condition | OrGroup] | None = None,
        cancel_if: list[Condition | OrGroup] | None = None,
    ) -> Callable[
        [Callable[Concatenate[TWorkflowInput, Context, P], R | CoroutineLike[R]]],
        Task[TWorkflowInput, R],
    ]:
        """
        A decorator to transform a function into a Hatchet task that runs as part of a workflow.

        :param name: The name of the task. If not specified, defaults to the name of the function being wrapped by the `task` decorator.

        :param schedule_timeout: The maximum time to wait for the task to be scheduled. The run will be canceled if the task does not begin within this time.

        :param execution_timeout: The maximum time to wait for the task to complete. The run will be canceled if the task does not complete within this time.

        :param parents: A list of tasks that are parents of the task. Note: Parents must be defined before their children.

        :param retries: The number of times to retry the task before failing.

        :param rate_limits: A list of rate limit configurations for the task.

        :param desired_worker_labels: A dictionary of desired worker labels that determine to which worker the task should be assigned. See documentation and examples on affinity and worker labels for more details.

        :param backoff_factor: The backoff factor for controlling exponential backoff in retries.

        :param backoff_max_seconds: The maximum number of seconds to allow retries with exponential backoff to continue.

        :param concurrency: A list of concurrency expressions for the task. If an integer is provided, it is treated as a constant concurrency limit with a `GROUP_ROUND_ROBIN` strategy, which means that only `N` runs of the task may execute at any given time.

        :param wait_for: A list of conditions that must be met before the task can run.

        :param skip_if: A list of conditions that, if met, will cause the task to be skipped.

        :param cancel_if: A list of conditions that, if met, will cause the task to be canceled.

        :returns: A decorator which creates a `Task` object.
        """

        computed_params = ComputedTaskParameters(
            schedule_timeout=schedule_timeout,
            execution_timeout=execution_timeout,
            retries=retries,
            backoff_factor=backoff_factor,
            backoff_max_seconds=backoff_max_seconds,
            task_defaults=self.config.task_defaults,
        )

        def inner(
            func: Callable[
                Concatenate[TWorkflowInput, Context, P], R | CoroutineLike[R]
            ],
        ) -> Task[TWorkflowInput, R]:
            task = Task(
                _fn=func,
                is_durable=False,
                workflow=self,
                type=StepType.DEFAULT,
                name=self._parse_task_name(name, func),
                execution_timeout=computed_params.execution_timeout,
                schedule_timeout=computed_params.schedule_timeout,
                parents=parents,
                retries=computed_params.retries,
                rate_limits=[r.to_proto() for r in rate_limits or []],
                desired_worker_labels={
                    key: transform_desired_worker_label(d)
                    for key, d in (desired_worker_labels or {}).items()
                },
                backoff_factor=computed_params.backoff_factor,
                backoff_max_seconds=computed_params.backoff_max_seconds,
                concurrency=concurrency,
                wait_for=wait_for,
                skip_if=skip_if,
                cancel_if=cancel_if,
            )

            self._default_tasks.append(task)

            return task

        return inner

    def durable_task(
        self,
        name: str | None = None,
        schedule_timeout: Duration = timedelta(minutes=5),
        execution_timeout: Duration = timedelta(seconds=60),
        parents: list[Task[TWorkflowInput, Any]] | None = None,
        retries: int = 0,
        rate_limits: list[RateLimit] | None = None,
        desired_worker_labels: dict[str, DesiredWorkerLabel] | None = None,
        backoff_factor: float | None = None,
        backoff_max_seconds: int | None = None,
        concurrency: int | list[ConcurrencyExpression] | None = None,
        wait_for: list[Condition | OrGroup] | None = None,
        skip_if: list[Condition | OrGroup] | None = None,
        cancel_if: list[Condition | OrGroup] | None = None,
    ) -> Callable[
        [
            Callable[
                Concatenate[TWorkflowInput, DurableContext, P], R | CoroutineLike[R]
            ]
        ],
        Task[TWorkflowInput, R],
    ]:
        """
        A decorator to transform a function into a durable Hatchet task that runs as part of a workflow.

        **IMPORTANT:** This decorator creates a _durable_ task, which works using Hatchet's durable execution capabilities. This is an advanced feature of Hatchet.

        See the Hatchet docs for more information on durable execution to decide if this is right for you.

        :param name: The name of the task. If not specified, defaults to the name of the function being wrapped by the `task` decorator.

        :param schedule_timeout: The maximum time to wait for the task to be scheduled. The run will be canceled if the task does not begin within this time.

        :param execution_timeout: The maximum time to wait for the task to complete. The run will be canceled if the task does not complete within this time.

        :param parents: A list of tasks that are parents of the task. Note: Parents must be defined before their children.

        :param retries: The number of times to retry the task before failing.

        :param rate_limits: A list of rate limit configurations for the task.

        :param desired_worker_labels: A dictionary of desired worker labels that determine to which worker the task should be assigned. See documentation and examples on affinity and worker labels for more details.

        :param backoff_factor: The backoff factor for controlling exponential backoff in retries.

        :param backoff_max_seconds: The maximum number of seconds to allow retries with exponential backoff to continue.

        :param concurrency: A list of concurrency expressions for the task. If an integer is provided, it is treated as a constant concurrency limit with a `GROUP_ROUND_ROBIN` strategy, which means that only `N` runs of the task may execute at any given time.

        :param wait_for: A list of conditions that must be met before the task can run.

        :param skip_if: A list of conditions that, if met, will cause the task to be skipped.

        :param cancel_if: A list of conditions that, if met, will cause the task to be canceled.

        :returns: A decorator which creates a `Task` object.
        """

        computed_params = ComputedTaskParameters(
            schedule_timeout=schedule_timeout,
            execution_timeout=execution_timeout,
            retries=retries,
            backoff_factor=backoff_factor,
            backoff_max_seconds=backoff_max_seconds,
            task_defaults=self.config.task_defaults,
        )

        def inner(
            func: Callable[
                Concatenate[TWorkflowInput, DurableContext, P], R | CoroutineLike[R]
            ],
        ) -> Task[TWorkflowInput, R]:
            task = Task(
                _fn=func,
                is_durable=True,
                workflow=self,
                type=StepType.DEFAULT,
                name=self._parse_task_name(name, func),
                execution_timeout=computed_params.execution_timeout,
                schedule_timeout=computed_params.schedule_timeout,
                parents=parents,
                retries=computed_params.retries,
                rate_limits=[r.to_proto() for r in rate_limits or []],
                desired_worker_labels={
                    key: transform_desired_worker_label(d)
                    for key, d in (desired_worker_labels or {}).items()
                },
                backoff_factor=computed_params.backoff_factor,
                backoff_max_seconds=computed_params.backoff_max_seconds,
                concurrency=concurrency,
                wait_for=wait_for,
                skip_if=skip_if,
                cancel_if=cancel_if,
            )

            self._durable_tasks.append(task)

            return task

        return inner

    def on_failure_task(
        self,
        name: str | None = None,
        schedule_timeout: Duration = timedelta(minutes=5),
        execution_timeout: Duration = timedelta(seconds=60),
        retries: int = 0,
        rate_limits: list[RateLimit] | None = None,
        backoff_factor: float | None = None,
        backoff_max_seconds: int | None = None,
        concurrency: int | list[ConcurrencyExpression] | None = None,
    ) -> Callable[
        [Callable[Concatenate[TWorkflowInput, Context, P], R | CoroutineLike[R]]],
        Task[TWorkflowInput, R],
    ]:
        """
        A decorator to transform a function into a Hatchet on-failure task that runs as the last step in a workflow that had at least one task fail.

        :param name: The name of the on-failure task. If not specified, defaults to the name of the function being wrapped by the `on_failure_task` decorator.

        :param schedule_timeout: The maximum time to wait for the task to be scheduled. The run will be canceled if the task does not begin within this time.

        :param execution_timeout: The maximum time to wait for the task to complete. The run will be canceled if the task does not complete within this time.

        :param retries: The number of times to retry the on-failure task before failing.

        :param rate_limits: A list of rate limit configurations for the on-failure task.

        :param backoff_factor: The backoff factor for controlling exponential backoff in retries.

        :param backoff_max_seconds: The maximum number of seconds to allow retries with exponential backoff to continue.

        :param concurrency: A list of concurrency expressions for the on-failure task. If an integer is provided, it is treated as a constant concurrency limit with a `GROUP_ROUND_ROBIN` strategy, which means that only `N` runs of the task may execute at any given time.

        :returns: A decorator which creates a `Task` object.
        """

        def inner(
            func: Callable[
                Concatenate[TWorkflowInput, Context, P], R | CoroutineLike[R]
            ],
        ) -> Task[TWorkflowInput, R]:
            task = Task(
                is_durable=False,
                _fn=func,
                workflow=self,
                type=StepType.ON_FAILURE,
                name=self._parse_task_name(name, func) + "-on-failure",
                execution_timeout=execution_timeout,
                schedule_timeout=schedule_timeout,
                retries=retries,
                rate_limits=[r.to_proto() for r in rate_limits or []],
                backoff_factor=backoff_factor,
                backoff_max_seconds=backoff_max_seconds,
                concurrency=concurrency,
                desired_worker_labels=None,
                parents=None,
                wait_for=None,
                skip_if=None,
                cancel_if=None,
            )

            if self._on_failure_task:
                raise ValueError("Only one on-failure task is allowed")

            self._on_failure_task = task

            return task

        return inner

    def on_success_task(
        self,
        name: str | None = None,
        schedule_timeout: Duration = timedelta(minutes=5),
        execution_timeout: Duration = timedelta(seconds=60),
        retries: int = 0,
        rate_limits: list[RateLimit] | None = None,
        backoff_factor: float | None = None,
        backoff_max_seconds: int | None = None,
        concurrency: int | list[ConcurrencyExpression] | None = None,
    ) -> Callable[
        [Callable[Concatenate[TWorkflowInput, Context, P], R | CoroutineLike[R]]],
        Task[TWorkflowInput, R],
    ]:
        """
        A decorator to transform a function into a Hatchet on-success task that runs as the last step in a workflow that had all upstream tasks succeed.

        :param name: The name of the on-success task. If not specified, defaults to the name of the function being wrapped by the `on_success_task` decorator.

        :param schedule_timeout: The maximum time to wait for the task to be scheduled. The run will be canceled if the task does not begin within this time.

        :param execution_timeout: The maximum time to wait for the task to complete. The run will be canceled if the task does not complete within this time.

        :param retries: The number of times to retry the on-success task before failing

        :param rate_limits: A list of rate limit configurations for the on-success task.

        :param backoff_factor: The backoff factor for controlling exponential backoff in retries.

        :param backoff_max_seconds: The maximum number of seconds to allow retries with exponential backoff to continue.

        :param concurrency: A list of concurrency expressions for the on-success task. If an integer is provided, it is treated as a constant concurrency limit with a `GROUP_ROUND_ROBIN` strategy, which means that only `N` runs of the task may execute at any given time.

        :returns: A decorator which creates a Task object.
        """

        def inner(
            func: Callable[
                Concatenate[TWorkflowInput, Context, P], R | CoroutineLike[R]
            ],
        ) -> Task[TWorkflowInput, R]:
            task = Task(
                is_durable=False,
                _fn=func,
                workflow=self,
                type=StepType.ON_SUCCESS,
                name=self._parse_task_name(name, func) + "-on-success",
                execution_timeout=execution_timeout,
                schedule_timeout=schedule_timeout,
                retries=retries,
                rate_limits=[r.to_proto() for r in rate_limits or []],
                backoff_factor=backoff_factor,
                backoff_max_seconds=backoff_max_seconds,
                concurrency=concurrency,
                parents=None,
                desired_worker_labels=None,
                wait_for=None,
                skip_if=None,
                cancel_if=None,
            )

            if self._on_success_task:
                raise ValueError("Only one on-success task is allowed")

            self._on_success_task = task

            return task

        return inner

    def add_task(self, task: "Standalone[TWorkflowInput, Any]") -> None:
        """
        Add a task to a workflow. Intended to be used with a previously existing task (a Standalone),
        such as one created with `@hatchet.task()`, which has been converted to a `Task` object using `to_task`.

        For example:

        ```python
        @hatchet.task()
        def my_task(input, ctx) -> None:
            pass

        wf = hatchet.workflow()

        wf.add_task(my_task.to_task())
        ```
        """
        _task = task._task

        match _task.type:
            case StepType.DEFAULT:
                self._default_tasks.append(_task)
            case StepType.ON_FAILURE:
                if self._on_failure_task:
                    raise ValueError("Only one on-failure task is allowed")

                self._on_failure_task = _task
            case StepType.ON_SUCCESS:
                if self._on_success_task:
                    raise ValueError("Only one on-success task is allowed")

                self._on_success_task = _task
            case _:
                raise ValueError("Invalid task type")


class TaskRunRef(Generic[TWorkflowInput, R]):
    def __init__(
        self,
        standalone: "Standalone[TWorkflowInput, R]",
        workflow_run_ref: WorkflowRunRef,
    ):
        self._s = standalone
        self._wrr = workflow_run_ref

        self.workflow_run_id = workflow_run_ref.workflow_run_id

    def __str__(self) -> str:
        return self.workflow_run_id

    async def aio_result(self) -> R:
        result = await self._wrr.workflow_run_listener.aio_result(
            self._wrr.workflow_run_id
        )
        return self._s._extract_result(result)

    def result(self) -> R:
        result = self._wrr.result()

        return self._s._extract_result(result)

    def stream(self) -> RunEventListener:
        return self._wrr.stream()


class Standalone(BaseWorkflow[TWorkflowInput], Generic[TWorkflowInput, R]):
    def __init__(
        self, workflow: Workflow[TWorkflowInput], task: Task[TWorkflowInput, R]
    ) -> None:
        super().__init__(config=workflow.config, client=workflow.client)

        ## NOTE: This is a hack to assign the task back to the base workflow,
        ## since the decorator to mutate the tasks is not being called.
        self._default_tasks = [task]

        self._workflow = workflow
        self._task = task

        return_type = get_type_hints(self._task.fn).get("return")

        self._output_validator: TypeAdapter[TaskPayloadForInternalUse] = TypeAdapter(
            normalize_validator(return_type)
        )

        self.config = self._workflow.config

    @overload
    def _extract_result(self, result: dict[str, Any]) -> R: ...

    @overload
    def _extract_result(self, result: BaseException) -> BaseException: ...

    def _extract_result(
        self, result: dict[str, Any] | BaseException
    ) -> R | BaseException:
        if isinstance(result, BaseException):
            return result

        ## if a task is cancelled, we can get `None` back here
        ## this is a bit of an edge case since both `None` and an empty dict
        ## would cause Pydantic validation errors, but if you were expecting a `dict`
        ## return, then the empty dict would not error and would work correctly
        output = result.get(self._task.name) or {}

        return cast(
            R,
            self._output_validator.validate_python(
                output, context=HATCHET_PYDANTIC_SENTINEL
            ),
        )

    def run(
        self,
        input: TWorkflowInput = cast(TWorkflowInput, EmptyModel()),
        options: TriggerWorkflowOptions = TriggerWorkflowOptions(),
    ) -> R:
        """
        Run the workflow synchronously and wait for it to complete.

        This method triggers a workflow run, blocks until completion, and returns the extracted result.

        :param input: The input data for the workflow.
        :param options: Additional options for workflow execution.

        :returns: The extracted result of the workflow execution.
        """
        return self._extract_result(self._workflow.run(input, options))

    async def aio_run(
        self,
        input: TWorkflowInput = cast(TWorkflowInput, EmptyModel()),
        options: TriggerWorkflowOptions = TriggerWorkflowOptions(),
    ) -> R:
        """
        Run the workflow asynchronously and wait for it to complete.

        This method triggers a workflow run, awaits until completion, and returns the extracted result.

        :param input: The input data for the workflow, must match the workflow's input type.
        :param options: Additional options for workflow execution like metadata and parent workflow ID.

        :returns: The extracted result of the workflow execution.
        """
        result = await self._workflow.aio_run(input, options)
        return self._extract_result(result)

    def run_no_wait(
        self,
        input: TWorkflowInput = cast(TWorkflowInput, EmptyModel()),
        options: TriggerWorkflowOptions = TriggerWorkflowOptions(),
    ) -> TaskRunRef[TWorkflowInput, R]:
        """
        Trigger a workflow run without waiting for it to complete.

        This method triggers a workflow run and immediately returns a reference to the run without blocking while the workflow runs.

        :param input: The input data for the workflow, must match the workflow's input type.
        :param options: Additional options for workflow execution like metadata and parent workflow ID.

        :returns: A `TaskRunRef` object representing the reference to the workflow run.
        """
        ref = self._workflow.run_no_wait(input, options)

        return TaskRunRef[TWorkflowInput, R](self, ref)

    async def aio_run_no_wait(
        self,
        input: TWorkflowInput = cast(TWorkflowInput, EmptyModel()),
        options: TriggerWorkflowOptions = TriggerWorkflowOptions(),
    ) -> TaskRunRef[TWorkflowInput, R]:
        """
        Asynchronously trigger a workflow run without waiting for it to complete.
        This method is useful for starting a workflow run and immediately returning a reference to the run without blocking while the workflow runs.

        :param input: The input data for the workflow.
        :param options: Additional options for workflow execution.

        :returns: A `TaskRunRef` object representing the reference to the workflow run.
        """
        ref = await self._workflow.aio_run_no_wait(input, options)

        return TaskRunRef[TWorkflowInput, R](self, ref)

    @overload
    def run_many(
        self,
        workflows: list[WorkflowRunTriggerConfig],
        return_exceptions: Literal[True],
    ) -> list[R | BaseException]: ...

    @overload
    def run_many(
        self,
        workflows: list[WorkflowRunTriggerConfig],
        return_exceptions: Literal[False] = False,
    ) -> list[R]: ...

    def run_many(
        self, workflows: list[WorkflowRunTriggerConfig], return_exceptions: bool = False
    ) -> list[R] | list[R | BaseException]:
        """
        Run a workflow in bulk and wait for all runs to complete.
        This method triggers multiple workflow runs, blocks until all of them complete, and returns the final results.

        :param workflows: A list of `WorkflowRunTriggerConfig` objects, each representing a workflow run to be triggered.
        :param return_exceptions: If `True`, exceptions will be returned as part of the results instead of raising them.
        :returns: A list of results for each workflow run.
        """
        return [
            self._extract_result(result)
            for result in self._workflow.run_many(
                workflows,
                ## hack: typing needs literal
                True if return_exceptions else False,  # noqa: SIM210
            )
        ]

    @overload
    async def aio_run_many(
        self,
        workflows: list[WorkflowRunTriggerConfig],
        return_exceptions: Literal[True],
    ) -> list[R | BaseException]: ...

    @overload
    async def aio_run_many(
        self,
        workflows: list[WorkflowRunTriggerConfig],
        return_exceptions: Literal[False] = False,
    ) -> list[R]: ...

    async def aio_run_many(
        self, workflows: list[WorkflowRunTriggerConfig], return_exceptions: bool = False
    ) -> list[R] | list[R | BaseException]:
        """
        Run a workflow in bulk and wait for all runs to complete.
        This method triggers multiple workflow runs, blocks until all of them complete, and returns the final results.

        :param workflows: A list of `WorkflowRunTriggerConfig` objects, each representing a workflow run to be triggered.
        :param return_exceptions: If `True`, exceptions will be returned as part of the results instead of raising them.
        :returns: A list of results for each workflow run.
        """
        return [
            self._extract_result(result)
            for result in await self._workflow.aio_run_many(
                workflows,
                ## hack: typing needs literal
                True if return_exceptions else False,  # noqa: SIM210
            )
        ]

    def run_many_no_wait(
        self, workflows: list[WorkflowRunTriggerConfig]
    ) -> list[TaskRunRef[TWorkflowInput, R]]:
        """
        Run a workflow in bulk without waiting for all runs to complete.

        This method triggers multiple workflow runs and immediately returns a list of references to the runs without blocking while the workflows run.

        :param workflows: A list of `WorkflowRunTriggerConfig` objects, each representing a workflow run to be triggered.
        :returns: A list of `WorkflowRunRef` objects, each representing a reference to a workflow run.
        """
        refs = self._workflow.run_many_no_wait(workflows)

        return [TaskRunRef[TWorkflowInput, R](self, ref) for ref in refs]

    async def aio_run_many_no_wait(
        self, workflows: list[WorkflowRunTriggerConfig]
    ) -> list[TaskRunRef[TWorkflowInput, R]]:
        """
        Run a workflow in bulk without waiting for all runs to complete.

        This method triggers multiple workflow runs and immediately returns a list of references to the runs without blocking while the workflows run.

        :param workflows: A list of `WorkflowRunTriggerConfig` objects, each representing a workflow run to be triggered.

        :returns: A list of `WorkflowRunRef` objects, each representing a reference to a workflow run.
        """
        refs = await self._workflow.aio_run_many_no_wait(workflows)

        return [TaskRunRef[TWorkflowInput, R](self, ref) for ref in refs]

    def mock_run(
        self,
        input: TWorkflowInput | None = None,
        additional_metadata: JSONSerializableMapping | None = None,
        parent_outputs: dict[str, JSONSerializableMapping] | None = None,
        retry_count: int = 0,
        lifespan: Any = None,
        dependencies: dict[str, Any] | None = None,
    ) -> R:
        """
        Mimic the execution of a task. This method is intended to be used to unit test
        tasks without needing to interact with the Hatchet engine. Use `mock_run` for sync
        tasks and `aio_mock_run` for async tasks.

        :param input: The input to the task.
        :param additional_metadata: Additional metadata to attach to the task.
        :param parent_outputs: Outputs from parent tasks, if any. This is useful for mimicking DAG functionality. For instance, if you have a task `step_2` that has a `parent` which is `step_1`, you can pass `parent_outputs={"step_1": {"result": "Hello, world!"}}` to `step_2.mock_run()` to be able to access `ctx.task_output(step_1)` in `step_2`.
        :param retry_count: The number of times the task has been retried.
        :param lifespan: The lifespan to be used in the task, which is useful if one was set on the worker. This will allow you to access `ctx.lifespan` inside of your task.
        :param dependencies: Dependencies to be injected into the task. This is useful for tasks that have dependencies defined using `Depends`. **IMPORTANT**: You must pass the dependencies _directly_, **not** the `Depends` objects themselves. For example, if you have a task that has a dependency `config: Annotated[str, Depends(get_config)]`, you should pass `dependencies={"config": "config_value"}` to `aio_mock_run`.

        :return: The output of the task.
        """

        return self._task.mock_run(
            input=input,
            additional_metadata=additional_metadata,
            parent_outputs=parent_outputs,
            retry_count=retry_count,
            lifespan=lifespan,
            dependencies=dependencies,
        )

    async def aio_mock_run(
        self,
        input: TWorkflowInput | None = None,
        additional_metadata: JSONSerializableMapping | None = None,
        parent_outputs: dict[str, JSONSerializableMapping] | None = None,
        retry_count: int = 0,
        lifespan: Any = None,
        dependencies: dict[str, Any] | None = None,
    ) -> R:
        """
        Mimic the execution of a task. This method is intended to be used to unit test
        tasks without needing to interact with the Hatchet engine. Use `mock_run` for sync
        tasks and `aio_mock_run` for async tasks.

        :param input: The input to the task.
        :param additional_metadata: Additional metadata to attach to the task.
        :param parent_outputs: Outputs from parent tasks, if any. This is useful for mimicking DAG functionality. For instance, if you have a task `step_2` that has a `parent` which is `step_1`, you can pass `parent_outputs={"step_1": {"result": "Hello, world!"}}` to `step_2.mock_run()` to be able to access `ctx.task_output(step_1)` in `step_2`.
        :param retry_count: The number of times the task has been retried.
        :param lifespan: The lifespan to be used in the task, which is useful if one was set on the worker. This will allow you to access `ctx.lifespan` inside of your task.
        :param dependencies: Dependencies to be injected into the task. This is useful for tasks that have dependencies defined using `Depends`. **IMPORTANT**: You must pass the dependencies _directly_, **not** the `Depends` objects themselves. For example, if you have a task that has a dependency `config: Annotated[str, Depends(get_config)]`, you should pass `dependencies={"config": "config_value"}` to `aio_mock_run`.

        :return: The output of the task.
        """

        return await self._task.aio_mock_run(
            input=input,
            additional_metadata=additional_metadata,
            parent_outputs=parent_outputs,
            retry_count=retry_count,
            lifespan=lifespan,
            dependencies=dependencies,
        )

    @property
    def is_async_function(self) -> bool:
        """
        Check if the task is an async function.

        :returns: True if the task is an async function, False otherwise.
        """
        return self._task.is_async_function

    def get_run_ref(self, run_id: str) -> TaskRunRef[TWorkflowInput, R]:
        """
        Get a reference to a task run by its run ID.

        :param run_id: The ID of the run to get the reference for.
        :returns: A `TaskRunRef` object representing the reference to the task run.
        """
        wrr = self._workflow.client._client.runs.get_run_ref(run_id)
        return TaskRunRef[TWorkflowInput, R](self, wrr)

    async def aio_get_result(self, run_id: str) -> R:
        """
        Get the result of a task run by its run ID.

        :param run_id: The ID of the run to get the result for.
        :returns: The result of the task run.
        """
        run_ref = self.get_run_ref(run_id)

        return await run_ref.aio_result()

    def get_result(self, run_id: str) -> R:
        """
        Get the result of a task run by its run ID.

        :param run_id: The ID of the run to get the result for.
        :returns: The result of the task run.
        """
        run_ref = self.get_run_ref(run_id)

        return run_ref.result()
