from hatchet_sdk import (
    Context,
    EmptyModel,
    Hatchet,
    StickyStrategy,
)

hatchet = Hatchet()

# > StickyWorker


sticky_workflow = hatchet.workflow(
    name="StickyWorkflow",
    # 👀 Specify a sticky strategy when declaring the workflow
    sticky=StickyStrategy.SOFT,
)


@sticky_workflow.task()
def step1a(input: EmptyModel, ctx: Context) -> dict[str, str | None]:
    return {"worker": ctx.worker_id}


@sticky_workflow.task()
def step1b(input: EmptyModel, ctx: Context) -> dict[str, str | None]:
    return {"worker": ctx.worker_id}



# > StickyChild

sticky_child_workflow = hatchet.workflow(
    name="StickyChildWorkflow", sticky=StickyStrategy.SOFT
)


@sticky_workflow.task(parents=[step1a, step1b])
async def step2(input: EmptyModel, ctx: Context) -> dict[str, str | None]:
    ref = await sticky_child_workflow.aio_run(
        sticky=True,
        wait_for_result=False,
    )

    await ref.aio_result()

    return {"worker": ctx.worker_id}


@sticky_child_workflow.task()
def child(input: EmptyModel, ctx: Context) -> dict[str, str | None]:
    return {"worker": ctx.worker_id}




def main() -> None:
    worker = hatchet.worker(
        "sticky-worker", slots=10, workflows=[sticky_workflow, sticky_child_workflow]
    )
    worker.start()


if __name__ == "__main__":
    main()
