{# Copyright 2021 Curtin University
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#   http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# Author: Tuan Chien -#}
import pendulum
from airflow.models.taskinstance import TaskInstance
from typing import Dict, List

from {{ package_name }}.config import schema_folder as default_schema_folder
from observatory.platform.workflows.stream_telescope import StreamRelease, StreamTelescope
from observatory.platform.utils.airflow_utils import AirflowVars


class {{ workflow_class }}Release(StreamRelease):
    def __init__(self, dag_id: str, start_date: pendulum.DateTime, end_date: pendulum.DateTime, first_release: bool):
        """Construct a {{ workflow_class }}Release instance

        :param dag_id: the id of the DAG.
        :param start_date: the start_date of the release.
        :param end_date: the end_date of the release.
        :param first_release: whether this is the first release that is processed for this DAG
        """

        super().__init__(dag_id, start_date, end_date, first_release)


class {{ workflow_class }}(StreamTelescope):
    """ {{ workflow_class}} Telescope."""

    DAG_ID = "{{ workflow_module }}"

    def __init__(
        self,
        workflow_id: int,
        dag_id: str = DAG_ID,
        start_date: pendulum.DateTime = pendulum.datetime(2020, 1, 1),
        schedule_interval: str = "@weekly",
        dataset_id: str = "your_dataset_id",
        dataset_description: str = "The your_dataset_name dataset: https://dataseturl",
        merge_partition_field: str = "id",
        dataset_type_id: str = "API dataset_type_id",
        batch_load: bool = True,
        load_bigquery_table_kwargs: Dict = None,
        table_descriptions: Dict = None,
        schema_folder: str = default_schema_folder(),
        airflow_vars: List = None,
        airflow_conns: List = None,
    ):
        """Construct a {{ workflow_class }} telescope instance.

        :param dag_id: the id of the DAG.
        :param start_date: the start date of the DAG.
        :param schedule_interval: the schedule interval of the DAG.
        :param dataset_id: the dataset id.
        :param workflow_id: API workflow id.
        :param dataset_description: the dataset description.
        :param merge_partition_field: the BigQuery field used to match partitions for a merge
        :param table_descriptions: a dictionary with table ids and corresponding table descriptions.
        :param schema_folder: the path to the SQL schema folder.
        :param batch_load: whether all files in the transform folder are loaded into 1 table at once
        :param airflow_vars: list of airflow variable keys, for each variable it is checked if it exists in airflow
        :param airflow_conns: list of airflow connection keys, for each connection it is checked if it exists in airflow
        :param dataset_type_id: API dataset_type_id.
        """

        if table_descriptions is None:
            table_descriptions = {dag_id: "Table with up to date data."}

        if airflow_vars is None:
            airflow_vars = [
                AirflowVars.DATA_PATH,
                AirflowVars.PROJECT_ID,
                AirflowVars.DATA_LOCATION,
                AirflowVars.DOWNLOAD_BUCKET,
                AirflowVars.TRANSFORM_BUCKET,
            ]

        # if airflow_conns is None:
        #     airflow_conns = [AirflowConns.SOMEDEFAULT_CONNECTION]

        super().__init__(
            dag_id,
            start_date,
            schedule_interval,
            dataset_id,
            merge_partition_field,
            schema_folder,
            batch_load=batch_load,
            load_bigquery_table_kwargs=load_bigquery_table_kwargs,
            dataset_description=dataset_description,
            table_descriptions=table_descriptions,
            airflow_vars=airflow_vars,
            airflow_conns=airflow_conns,
            workflow_id=workflow_id,
            dataset_type_id=dataset_type_id
        )

        # Add sensor tasks
        # self.add_operator(some_airflow_sensor)

        # Add setup tasks
        self.add_setup_task(self.check_dependencies)

        # Add ETL tasks
        self.add_task(self.task1)  # User provided
        # self.add_task(self.upload_transformed)  # From StreamTelescope

        # BQ loading functions from StreamTelescope
        # self.add_task(self.bq_load_partition)
        # self.add_task_chain([self.bq_delete_old,
        #                     self.bq_append_new], trigger_rule="none_failed")

        # cleanup
        self.add_task(self.cleanup)  # From StreamTelescope

    def make_release(self, **kwargs) -> {{ workflow_class }}Release:
        """Make a Release instance

        :param kwargs: The context passed from the PythonOperator.
        :return: {{ workflow_class }}Release
        """

        start_date, end_date, first_release = self.get_release_info(**kwargs)

        return {{ workflow_class }}Release(self.dag_id, start_date, end_date, first_release)

    def task1(self, release: {{ workflow_class }}Release, **kwargs):
        """Add your own comments.

        :param release: A {{ workflow_class }} instance
        :param kwargs: The context passed from the PythonOperator.
        :return: None.
        """
        pass
