from collections import namedtuple
from os.path import dirname, join, realpath
from urllib.parse import urlparse

import pendulum
import yaml
from airflow.datasets import Dataset
from airflow.decorators import dag, task, task_group
from airflow.operators.python import get_current_context
from airflow.providers.mysql.hooks.mysql import MySqlHook

Table = namedtuple("Table", ["conn_id", "database", "table"])


def parse_uri(uri: str) -> Table:
    o = urlparse(uri)
    if o.scheme != "mysql":
        raise ValueError("only mysql dataset is supported")
    ps = o.path.strip("/").split("/")
    return Table(o.netloc, ps[0], ps[1])


@dag(
    start_date=pendulum.datetime(2024, 11, 14, 7, 30, tz="UTC"),
    schedule="* * * * *",
    catchup=False,
    max_active_runs=1,
)
def auto_trigger():
    dataset_file = realpath(join(dirname(__file__), "datasets.yaml"))
    with open(dataset_file) as f:
        datasets = yaml.safe_load(f)

    for dataset in datasets:

        @task.short_circuit
        def poll(uri: str):
            context = get_current_context()
            start = context["data_interval_start"]
            tb = parse_uri(uri)
            db = MySqlHook(mysql_conn_id=tb.conn_id)
            return (
                db.get_pandas_df(
                    f"""
                    select 1 from {tb.database}.{tb.table} 
                    where created_at > %s
                    """,
                    (start.isoformat(),),
                ).size
                > 0
            )

        @task(outlets=[Dataset(dataset["uri"])])
        def trigger():
            # do nothing
            pass

        @task_group(group_id=dataset["name"])
        def group():
            poll(dataset["uri"]) >> trigger()

        group()


auto_trigger()