import asyncio
from pathlib import Path

from prefect import flow, get_client
from prefect.client.schemas.actions import WorkPoolCreate
from prefect.deployments import run_deployment
from prefect.variables import Variable


@flow
def dummy_flow():
    print("I am a dummy flow")


async def create_work_pool():
    async with get_client() as client:
        await client.create_work_pool(
            WorkPoolCreate(
                name="ecs-work-pool",
                type="ecs",
                description="A work pool for ECS work",
                base_job_template=dict(
                    job_configuration=dict(vpc_id="{{ vpc_id }}"),
                    variables=dict(
                        type="object",
                        properties=dict(
                            vpc_id=dict(
                                type="string",
                                title="VPC ID",
                                description="The VPC ID to use for this work pool",
                            )
                        ),
                    ),
                ),
            ),
            overwrite=True,
        )


if __name__ == "__main__":
    Variable.set("some_vpc_id", "vpc-12345678", overwrite=True)
    asyncio.run(create_work_pool())
    dummy_flow.from_source(
        source=str((p := Path(__file__)).parent.resolve()),
        entrypoint=f"{p.name}:dummy_flow",
    ).deploy(name="dummy", work_pool_name="ecs-work-pool")
    flow_run = run_deployment(
        "dummy-flow/dummy",
        timeout=0,
        job_variables=dict(vpc_id=Variable.get("some_vpc_id")),
    )
    assert (
        flow_run.job_variables.get("vpc_id") == "vpc-12345678"
    ), flow_run.job_variables