from datetime import timedelta
from typing import Any

from pydantic import BaseModel

from hatchet_sdk import Context, Hatchet

hatchet = Hatchet()


# > FanoutParent
class ParentInput(BaseModel):
    n: int = 100


class ChildInput(BaseModel):
    a: str


parent_wf = hatchet.workflow(name="FanoutParent", input_validator=ParentInput)
child_wf = hatchet.workflow(name="FanoutChild", input_validator=ChildInput)


@parent_wf.task(execution_timeout=timedelta(minutes=5))
async def spawn(input: ParentInput, ctx: Context) -> dict[str, Any]:
    print("spawning child")

    result = await child_wf.aio_run_many(
        [
            child_wf.create_bulk_run_item(
                input=ChildInput(a=str(i)),
                additional_metadata={"hello": "earth"},
                key=f"child{i}",
            )
            for i in range(input.n)
        ],
    )

    print(f"results {result}")

    return {"results": result}


# !!


# > FanoutChild
@child_wf.task()
async def process(input: ChildInput, ctx: Context) -> dict[str, str]:
    print(f"child process {input.a}")
    return {"status": input.a}


@child_wf.task(parents=[process])
async def process2(input: ChildInput, ctx: Context) -> dict[str, str]:
    process_output = ctx.task_output(process)
    a = process_output["status"]

    return {"status2": a + "2"}


# !!


def main() -> None:
    worker = hatchet.worker("fanout-worker", slots=40, workflows=[parent_wf, child_wf])
    worker.start()


if __name__ == "__main__":
    main()
