from datetime import timedelta, datetime
from typing import Union, Iterable, Dict

from airflow.operators.branch import BaseBranchOperator
from airflow.utils import timezone


class BackfillDetectorBranchOperator(BaseBranchOperator):
    def __init__(
            self,
            go_to_task_if_true: Union[str, Iterable[str]],
            go_to_task_if_false: Union[str, Iterable[str]],
            threshold: timedelta,
            **kwargs,
    ) -> None:
        super().__init__(**kwargs)
        self.threshold = threshold
        self.go_to_task_if_true = go_to_task_if_true
        self.follow_task_ids_if_false = go_to_task_if_false

    def choose_branch(self, context: Dict) -> Union[str, Iterable[str]]:
        data_interval_end = timezone.make_naive(context["data_interval_end"], self.dag.timezone)
        this_dag_run_threshold = data_interval_end + self.threshold
        self.log.info(f'This dag_run have threshold up to "{this_dag_run_threshold}"')

        if datetime.now() > this_dag_run_threshold:
            self.log.info('This run considered as backfill')
            return self.go_to_task_if_true
        else:
            self.log.info('This run considered as NOT backfill')
            return self.follow_task_ids_if_false