import prefect
from datetime import datetime
from pytz import timezone


from prefect import task, Flow
from prefect.environments import LocalEnvironment
from prefect.engine.executors import DaskExecutor
from prefect.environments.storage import Docker

from prefect.schedules import Schedule
from prefect.schedules.clocks import CronClock
from prefect.tasks.prefect import RenameFlowRun


class MyRenameTask(RenameFlowRun):
    def run(self):
        logger = prefect.context.get("logger")
        logger.info(prefect.context.flow_run_id)
        logger.info(f"run_{(prefect.context.scheduled_start_time).strftime('%Y-%m-%d')}")
        return super().run(
            flow_run_id=prefect.context.flow_run_id,
            flow_run_name=f"run_{(prefect.context.scheduled_start_time).strftime('%Y-%m-%d')}",
        )


tz = timezone("Europe/Prague")
schedule = Schedule(clocks=[CronClock("0 0 * * *", start_date=datetime.now(tz=tz))])


@task(
    name="Task name",
    slug="Task slug",
    task_run_name=lambda **kwargs: kwargs["scheduled_start_time"].strftime("%Y-%m-%d"),
)
def say_hello():
    logger = prefect.context.get("logger")
    logger.info("Hello, world!")


@task
def rename_flow_run():
    MyRenameTask().run()


with Flow(
    "Test flow (aka DAG)",
    environment=LocalEnvironment(
        executor=DaskExecutor(
            "tcp://***:8786",
        )
    ),
) as f:
    rename_flow_run(upstream_tasks=[say_hello])

    f.schedule = schedule

    # Something to get flow_run_id ???
    # from prefect import Client
    # client = Client()
    # client.set_flow_run_name(flow_run_id, flow_run_name)

    f.storage = Docker(
        registry_url="***/lmc-bi-prefect-flows",
        # python_dependencies=["pyqlikengine"],
    )

    f.register(project_name="Test")
    # f.run()