# Copyright 2020 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, James Diprose, Tuan Chien

import json
import logging
import os
import pathlib
from typing import Dict, List, Tuple, Optional

import pendulum
from airflow.exceptions import AirflowException
from airflow.models import Variable
from google.cloud.bigquery import SourceFormat

from observatory.platform.utils.airflow_utils import AirflowVars
from observatory.platform.utils.gc_utils import bigquery_sharded_table_id, create_bigquery_snapshot
from observatory.platform.utils.release_utils import (
    get_dataset_releases,
    get_datasets,
    get_latest_dataset_release,
    is_first_release,
)
from observatory.platform.utils.workflow_utils import (
    batch_blob_name,
    blob_name,
    bq_append_from_file,
    bq_append_from_partition,
    bq_delete_old,
    bq_load_ingestion_partition,
    delete_old_xcoms,
    table_ids_from_path,
    upload_files_from_list,
)
from observatory.platform.workflows.workflow import Release, Workflow


class StreamRelease(Release):
    def __init__(
        self,
        dag_id: str,
        start_date: pendulum.DateTime,
        end_date: pendulum.DateTime,
        first_release: bool,
        download_files_regex: str = None,
        extract_files_regex: str = None,
        transform_files_regex: str = None,
    ):
        """Construct a StreamRelease 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
        :param download_files_regex: regex pattern that is used to find files in download folder
        :param extract_files_regex: regex pattern that is used to find files in extract folder
        :param transform_files_regex: regex pattern that is used to find files in transform folder
        """
        self.start_date = start_date
        self.end_date = end_date
        self.first_release = first_release
        release_id = self.start_date.strftime("%Y_%m_%d") + "-" + self.end_date.strftime("%Y_%m_%d")
        super().__init__(dag_id, release_id, download_files_regex, extract_files_regex, transform_files_regex)


class StreamTelescope(Workflow):
    def __init__(
        self,
        dag_id: str,
        start_date: pendulum.DateTime,
        schedule_interval: str,
        dataset_id: str,
        merge_partition_field: str,
        schema_folder: str,
        workflow_id: int,
        dataset_type_id: str,
        queue: str = "default",
        max_retries: int = 3,
        max_active_runs: int = 1,
        source_format: str = SourceFormat.NEWLINE_DELIMITED_JSON,
        schema_prefix: str = "",
        schema_version: str = None,
        load_bigquery_table_kwargs: Dict = None,
        dataset_description: str = "",
        table_descriptions: Dict[str, str] = None,
        batch_load: bool = False,
        airflow_vars: list = None,
        airflow_conns: list = None,
        tags: Optional[List[str]] = None,
        **kwargs,
    ):
        """Construct a StreamTelescope 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 merge_partition_field: the BigQuery field used to match partitions for a merge
        :param schema_folder: the path to the SQL schema folder.
        :param workflow_id: Observatory API workflow id.
        :param dataset_type_id: Observatory API type_id for the top level dataset type.
        :param queue: the Airflow queue name.
        :param max_retries: the number of times to retry each task.
        :param max_active_runs: the maximum number of DAG runs that can be run at once.
        :param source_format: the format of the data to load into BigQuery.
        :param schema_prefix: the prefix used to find the schema path
        :param schema_version: the version used to find the schema path
        :param load_bigquery_table_kwargs: the customisation parameters for loading data into a BigQuery table.
        :param dataset_description: description for the BigQuery dataset.
        :param table_descriptions: a dictionary with table ids and corresponding table descriptions
        :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 tags: Optional Airflow DAG tags to add.
        """

        # Set transform_bucket_name as required airflow variable
        if not airflow_vars:
            airflow_vars = []
        airflow_vars = list(set([AirflowVars.TRANSFORM_BUCKET] + airflow_vars))

        super().__init__(
            dag_id,
            start_date,
            schedule_interval,
            catchup=False,
            queue=queue,
            max_retries=max_retries,
            max_active_runs=max_active_runs,
            airflow_vars=airflow_vars,
            airflow_conns=airflow_conns,
            workflow_id=workflow_id,
            dataset_type_id=dataset_type_id,
            tags=tags,
            **kwargs,
        )

        self.schema_prefix = schema_prefix
        self.schema_version = schema_version
        self.dataset_id = dataset_id
        self.source_format = source_format
        self.merge_partition_field = merge_partition_field
        self.schema_folder = schema_folder
        self.load_bigquery_table_kwargs = load_bigquery_table_kwargs if load_bigquery_table_kwargs else dict()
        self.dataset_description = dataset_description
        self.table_descriptions = table_descriptions if table_descriptions else dict()
        self.batch_load = batch_load

    def get_release_info(self, **kwargs) -> Tuple[pendulum.DateTime, pendulum.DateTime, bool]:
        """Return release information with the start and end date. [start date, end date) Includes start date, excludes end date.
        :param kwargs: The context passed from the PythonOperator.
        :return: None.
        """

        # Check if first release or not
        first_release = is_first_release(self.workflow_id)
        if first_release:
            # When first release, set start date to the start of the DAG
            # Can we get away with using data_interval_start instead?
            start_date = pendulum.instance(kwargs["dag"].default_args["start_date"]).start_of("day")
        else:
            datasets = get_datasets(workflow_id=self.workflow_id)
            try:
                top_level_dataset = [
                    dataset for dataset in datasets if dataset.dataset_type.type_id == self.dataset_type_id
                ][0]
            except IndexError:
                raise AirflowException(
                    f"No top level dataset type set, expecting dataset type with type_id: {self.dataset_type_id}"
                )
            releases = get_dataset_releases(dataset_id=top_level_dataset.id)
            latest_release = get_latest_dataset_release(releases=releases)
            start_date = pendulum.instance(latest_release.end_date).start_of("day")

        end_date = pendulum.instance(kwargs["data_interval_end"])
        logging.info(f"Start date: {start_date}, end date: {end_date}, first release: {first_release}")

        return start_date, end_date, first_release

    def get_bq_load_info(self, release: StreamRelease) -> List[Tuple[str, str, str]]:
        """Get (a list of) the transform blob, main table id and partition table id that are used to load data into
        BigQuery.
        When the batch_load property of the telescope is set to True, all blobs in the transform folder will be loaded
        into 1 table at once.
        When the batch_load property of the telescope is set to False, each blob in the transform folder will be used to
        create/update an individual table.

        :param release: The release object.
        :return: List with tuples of transform_blob, main_table_id and partition_table_id
        """
        bq_info = []
        if self.batch_load:
            if release.transform_files:
                transform_blob = batch_blob_name(release.transform_folder)
                main_table_id, partition_table_id = (self.dag_id, f"{self.dag_id}_partitions")
                bq_info = [(transform_blob, main_table_id, partition_table_id)]
        else:
            for transform_path in release.transform_files:
                transform_blob = blob_name(transform_path)
                main_table_id, partition_table_id = table_ids_from_path(transform_path)
                bq_info.append((transform_blob, main_table_id, partition_table_id))
        bq_info.sort(key=lambda x: x[0])
        return bq_info

    def download(self, release: StreamRelease, **kwargs):
        """Task to download the StreamRelease release.

        :param release: a StreamRelease instance.
        :param kwargs: The context passed from the PythonOperator.
        :return: None.
        """
        release.download()

    def upload_downloaded(self, release: StreamRelease, **kwargs):
        """Upload the downloaded files for the given release.

        :param release: a StreamRelease instance
        :param kwargs: The context passed from the PythonOperator.
        :return: None.
        """
        upload_files_from_list(release.download_files, release.download_bucket)

    def extract(self, release: StreamRelease, **kwargs):
        """Task to extract the StreamRelease release.

        :param release: a StreamRelease instance.
        :param kwargs: The context passed from the PythonOperator.
        :return: None.
        """
        release.extract()

    def transform(self, release: StreamRelease, **kwargs):
        """Task to transform the StreamRelease release.

        :param release: a StreamRelease instance.
        :param kwargs: The context passed from the PythonOperator.
        :return: None.
        """
        release.transform()

    def upload_transformed(self, release: StreamRelease, **kwargs):
        """Upload the transformed files for the given release.

        :param release: a StreamRelease instance
        :param kwargs: The context passed from the PythonOperator.
        :return: None.
        """
        upload_files_from_list(release.transform_files, release.transform_bucket)

    def bq_load_partition(self, release: StreamRelease, **kwargs):
        """For each file listed in transform_files, create a time based partition in BigQuery

        :param release: a StreamRelease instance
        :param kwargs: The context passed from the PythonOperator.
        :return: None.
        """
        if release.first_release:
            logging.info("Skipped, because first release")
            return

        bq_load_info = self.get_bq_load_info(release)
        for transform_blob, main_table_id, partition_table_id in bq_load_info:
            table_description = self.table_descriptions.get(main_table_id, "")
            bq_load_ingestion_partition(
                self.schema_folder,
                release.end_date,
                transform_blob,
                self.dataset_id,
                main_table_id,
                partition_table_id,
                self.source_format,
                self.schema_prefix,
                self.schema_version,
                self.dataset_description,
                table_description=table_description,
                **self.load_bigquery_table_kwargs,
            )

    def bq_delete_old(self, release: StreamRelease, **kwargs):
        """Delete old rows from the 'main' table, based on rows that are in a partition of the 'partitions' table.

        :param release: a StreamRelease instance
        :param kwargs: The context passed from the PythonOperator.
        :return: None.
        """
        if release.first_release:
            logging.info("Skipped, because first release")
            return

        logging.info(f"Deleting old data from main table using partition with ingestion date of {release.end_date}")
        bytes_budget = kwargs.get("bytes_budget", None)
        bq_load_info = self.get_bq_load_info(release)
        project_id = Variable.get(AirflowVars.PROJECT_ID)
        for _, main_table_id, partition_table_id in bq_load_info:
            bq_delete_old(
                release.end_date,
                self.dataset_id,
                main_table_id,
                partition_table_id,
                self.merge_partition_field,
                bytes_budget=bytes_budget,
                project_id=project_id,
            )

    def bq_append_new(self, release: StreamRelease, **kwargs):
        """Append rows to the 'main' table, based on rows that are in a partition of the 'partitions' table.

        :param release: a StreamRelease instance
        :param kwargs: The context passed from the PythonOperator.
        :return: None.
        """
        bq_load_info = self.get_bq_load_info(release)

        if release.first_release:
            for transform_blob, main_table_id, partition_table_id in bq_load_info:
                table_description = self.table_descriptions.get(main_table_id, "")
                bq_append_from_file(
                    self.schema_folder,
                    release.end_date,
                    transform_blob,
                    self.dataset_id,
                    main_table_id,
                    self.source_format,
                    self.schema_prefix,
                    self.schema_version,
                    self.dataset_description,
                    table_description=table_description,
                    **self.load_bigquery_table_kwargs,
                )
            return

        logging.info(f"Appending data to main table from partition with ingestion date of {release.end_date}")
        for transform_blob, main_table_id, partition_table_id in bq_load_info:
            bq_append_from_partition(
                release.end_date,
                self.dataset_id,
                main_table_id,
                partition_table_id,
            )

    def bq_create_snapshot(self, release: StreamRelease, **kwargs):
        """Create a snapshot from an existing BigQuery table

        :param release: a StreamRelease instance
        :param kwargs: The context passed from the PythonOperator.
        :return: None.
        """
        logging.info("Creating a BigQuery snapshot")
        bq_load_info = self.get_bq_load_info(release)
        project_id = Variable.get(AirflowVars.PROJECT_ID)
        for _, main_table_id, _ in bq_load_info:
            destination_table_id = bigquery_sharded_table_id(main_table_id + "_snapshots", release.end_date)
            create_bigquery_snapshot(project_id, self.dataset_id, main_table_id, self.dataset_id, destination_table_id)

    def cleanup(self, release: StreamRelease, **kwargs):
        """Delete downloaded, extracted and transformed files of the release. Deletes old xcoms.

        :param release: a StreamRelease instance
        :param kwargs: The context passed from the PythonOperator.
        :return: None.
        """
        release.cleanup()

        execution_date = kwargs["execution_date"]
        delete_old_xcoms(dag_id=self.dag_id, execution_date=execution_date)

    def merge_update_files(self, release: StreamRelease, **kwargs):
        """Merge the transformed jsonl into a single jsonl and delete the parts.

        :param release: a StreamRelease instance
        :param kwargs: The context passed from the PythonOperator.
        :return: None.
        """

        # Skip on first release.
        if release.first_release:
            return

        key = self.merge_partition_field
        input_files = release.transform_files
        input_files.sort()  # Assume lexicological order is correct order to apply updates.
        output_file = os.path.join(release.transform_folder, f"{self.DAG_ID}.jsonl")

        # Create mapping
        update_data = {}
        for i, input_file in enumerate(input_files):
            with open(input_file, "r") as f:
                for line in f:
                    row = json.loads(line)
                    update_data[row[key]] = i

        # Write out merged file
        with open(output_file, "w") as outf:
            for i, input_file in enumerate(input_files):
                with open(input_file, "r") as inf:
                    for line in inf:
                        row = json.loads(line)
                        if update_data[row[key]] == i:
                            outf.write(line)

        # Delete original parts
        for file in input_files:
            pathlib.Path(file).unlink()

        # Upload merged file to bucket.
        # Ideally in own task, but would it need branching mechanic (currently not implemented in Workflow).
        upload_files_from_list(release.transform_files, release.transform_bucket)
