import prefect
from prefect import Flow, task
from prefect.storage import Docker
from prefect.run_configs import KubernetesRun
from prefect.executors import DaskExecutor

from dask_kubernetes import KubeCluster, make_pod_spec

PROJECT_NAME = ...
FLOW_NAME = ...
IMAGE_TAG = ...

storage = Docker(
    registry_url = f"<http://eu.gcr.io/{PROJECT_NAME}|eu.gcr.io/{PROJECT_NAME}>",
    image_name = FLOW_NAME,
    image_tag  = IMAGE_TAG,
    dockerfile = "Dockerfile"
)

run_config = KubernetesRun(
    image = f"<http://eu.gcr.io/{PROJECT_NAME}/{FLOW_NAME}:{IMAGE_TAG}|eu.gcr.io/{PROJECT_NAME}/{FLOW_NAME}:{IMAGE_TAG}>", 
    labels = ["dask"])

executor = DaskExecutor(
        cluster_class=lambda: KubeCluster(
            make_pod_spec(
                image=prefect.context.image,
                extra_pod_config= {
                    "nodeSelector": {"<http://cloud.google.com/gke-nodepool|cloud.google.com/gke-nodepool>": "worker-pool"}
                }
        )),
        adapt_kwargs={"minimum": 1, "maximum": 3},
    )

with Flow(
    name = FLOW_NAME,    
    storage = storage,
    executor = executor,
    run_config = run_config) as flow:

    #tasks here