{# 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: Aniek Roelofs -#}
import logging
from typing import Dict, List, Optional

import pendulum
from airflow.exceptions import AirflowException

from observatory.api.client.model.organisation import Organisation
from {{ package_name }}.config import schema_folder as default_schema_folder
from observatory.platform.workflows.organisation_telescope import OrganisationRelease, OrganisationTelescope
from observatory.platform.utils.airflow_utils import AirflowConns, AirflowVars
from observatory.platform.utils.workflow_utils import make_dag_id


class {{ workflow_class }}Release(OrganisationRelease):
    def __init__(self, dag_id: str, release_date: pendulum.DateTime, organisation: Organisation):
        """Construct a {{ workflow_class }}Release.

        :param dag_id: the id of the DAG.
        :param release_date: the date of the release.
        :param organisation: the Organisation of which data is processed.
        """
        super().__init__(dag_id=dag_id, release_date=release_date, organisation=organisation)


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

    DAG_ID_PREFIX = "{{ workflow_module }}"

    def __init__(
        self,
        organisation: Organisation,
        extra1: str,
        dag_id: Optional[str] = None,
        start_date: pendulum.DateTime = pendulum.datetime(2020, 1, 1),
        schedule_interval: str = "@weekly",
        catchup: bool = True,
        dataset_id: str = "your_dataset_id",
        dataset_description: str = "The your_dataset dataset: https://dataseturl/",
        load_bigquery_table_kwargs: Dict = None,
        table_descriptions: Dict = None,
        schema_prefix: str = "",
        schema_folder: str = default_schema_folder(),
        airflow_vars=None,
        airflow_conns=None,
    ):
        """Construct a {{ workflow_class }}Telescope instance.

        :param organisation: the Organisation of which data is processed.
        :param extra1: the value for extra1, obtained from the 'extra' info from the API regarding the telescope.
        :param dag_id: the id of the DAG, by default this is automatically generated based on the DAG_ID_PREFIX and the
        organisation name.
        :param start_date: the start date of the DAG.
        :param schedule_interval: the schedule interval of the DAG.
        :param catchup: whether to catchup the DAG or not.
        :param dataset_id: the BigQuery dataset id.
        :param dataset_description: description for the BigQuery dataset.
        :param load_bigquery_table_kwargs: the customisation parameters for loading data into a BigQuery table.
        :param table_descriptions: a dictionary with table ids and corresponding table descriptions.
        :param schema_prefix: the prefix used to find the schema path.
        :param schema_folder: the path to the SQL schema folder.
        :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
        """
        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]

        if dag_id is None:
            dag_id = make_dag_id(self.DAG_ID_PREFIX, organisation.name)

        super().__init__(
            organisation,
            dag_id,
            start_date,
            schedule_interval,
            dataset_id,
            schema_folder,
            load_bigquery_table_kwargs=load_bigquery_table_kwargs,
            dataset_description=dataset_description,
            table_descriptions=table_descriptions,
            catchup=catchup,
            schema_prefix=schema_prefix,
            airflow_vars=airflow_vars,
            airflow_conns=airflow_conns,
        )

        self.extra1 = extra1

        # Add sensor tasks
        # self.add_operator(some_airflow_sensor)

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

        # Add ETL tasks
        self.add_task(self.task1)
        # self.add_task(self.upload_downloaded)  # From OrganisationTelescope
        # self.add_task(self.upload_transformed)  # From OrganisationTelescope
        # self.add_task(self.bq_load_partition)  # From OrganisationTelescope
        self.add_task(self.cleanup)  # From OrganisationTelescope

    def make_release(self, **kwargs) -> List[{{ workflow_class }}Release]:
        """Make release instances.

        :param kwargs: the context passed from the PythonOperator.
        :return: a list of {{ workflow_class }}Release instances.
        """
        release_date = kwargs["execution_date"]
        logging.info(f"Release date: {release_date}")
        releases = [{{ workflow_class }}Release(self.dag_id, release_date, self.organisation)]
        return releases

    def check_dependencies(self, **kwargs) -> bool:
        """Check dependencies of DAG. Add to parent method to additionally check for a expected extra key

        :return: True if dependencies are valid.
        """
        super().check_dependencies()

        if self.extra1 is None:
            expected_extra = {"extra1": "expected_value"}
            raise AirflowException(
                f"Value for extra1 is not set in 'extra' of telescope, " f"extra example: {expected_extra}"
            )
        return True
    
    def task1(self, releases: List[{{ workflow_class }}Release], **kwargs):
        """Add your own comments.

        :param releases: A list of {{ workflow_class }}Release instances
        :param kwargs: The context passed from the PythonOperator.
        :return: None.
        """
        pass
