# 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: James Diprose

import glob
import json
import logging
import os
import shutil
import subprocess
from dataclasses import dataclass
from typing import List

import pendulum
from academic_observatory_workflows.config import schema_folder
from airflow.exceptions import AirflowException
from airflow.models.taskinstance import TaskInstance
from airflow.models.variable import Variable
from google.cloud.bigquery import SourceFormat
from observatory.platform.utils.airflow_utils import AirflowVars, check_variables
from observatory.platform.utils.config_utils import find_schema
from observatory.platform.utils.gc_utils import (
    bigquery_sharded_table_id,
    bigquery_table_exists,
    create_bigquery_dataset,
    load_bigquery_table,
    upload_files_to_cloud_storage,
)
from observatory.platform.utils.http_download import download_file
from observatory.platform.utils.proc_utils import wait_for_process
from observatory.platform.utils.url_utils import retry_session
from observatory.platform.utils.workflow_utils import SubFolder, workflow_path

OPEN_CITATIONS_ARTICLE_ID = 6741422
OPEN_CITATIONS_VERSION_URL = "https://api.figshare.com/v2/articles/{article_id}/versions"


@dataclass
class File:
    name: str
    download_url: str
    md5hash: str
    parent: "OpenCitationsRelease" = None

    @property
    def download_blob_name(self):
        return f"telescopes/{OpenCitationsTelescope.DAG_ID}/{self.parent.release_name}/{self.name}"


class OpenCitationsRelease:
    release_date: pendulum.DateTime
    files: List[File]

    def __init__(self, release_date: pendulum.DateTime, files: List[File]):
        self.release_date = release_date
        self.files = files

        # Add parent release to files
        for file in self.files:
            file.parent = self

    @property
    def release_name(self) -> str:
        return f'{OpenCitationsTelescope.DAG_ID}_{self.release_date.strftime("%Y_%m_%d")}'

    @property
    def download_path(self) -> str:
        return os.path.join(workflow_path(SubFolder.downloaded, OpenCitationsTelescope.DAG_ID), self.release_name)

    @property
    def extract_path(self) -> str:
        return os.path.join(workflow_path(SubFolder.extracted, OpenCitationsTelescope.DAG_ID), self.release_name)

    @property
    def transformed_blob_path(self) -> str:
        return f"telescopes/{OpenCitationsTelescope.DAG_ID}/{self.release_name}/*.csv"


def list_open_citations_releases(
    start_date: pendulum.DateTime = None, end_date: pendulum.DateTime = None, timeout: float = 30.0
) -> List[OpenCitationsRelease]:
    """List the Open Citations releases for a given time period.

    :param start_date: the start date.
    :param end_date: the end date.
    :param timeout: timeout in seconds.
    :return: a list of Open Citations releases found within the given time period (if the time period was specified).
    """

    versions = fetch_open_citations_versions()
    releases = []
    for version in versions:
        version_url = version["url"]
        response = retry_session().get(version_url, timeout=timeout, headers={"Accept-encoding": "gzip"})
        article = json.loads(response.text)
        release_date = pendulum.parse(article["created_date"])

        if (start_date is None or start_date <= release_date) and (end_date is None or release_date <= end_date):
            files = []
            for file in article["files"]:
                name = file["name"]
                download_url = file["download_url"]
                md5hash = file["computed_md5"]
                files.append(File(name, download_url, md5hash))
            releases.append(OpenCitationsRelease(release_date, files))
    return releases


def fetch_open_citations_versions():
    """Fetch the Open Citation versions that are available.

    :return: the open citations versions.
    """

    url = OPEN_CITATIONS_VERSION_URL.format(article_id=OPEN_CITATIONS_ARTICLE_ID)
    response = retry_session().get(url, timeout=30, headers={"Accept-encoding": "gzip"})
    return json.loads(response.text)


def pull_releases(ti: TaskInstance) -> List[OpenCitationsRelease]:
    return ti.xcom_pull(
        key=OpenCitationsTelescope.RELEASES_TOPIC_NAME,
        task_ids=OpenCitationsTelescope.TASK_ID_LIST_RELEASES,
        include_prior_dates=False,
    )


class OpenCitationsTelescope:
    """A container for holding the constants and static functions for the Open Citations telescope."""

    DAG_ID = "open_citations"
    DESCRIPTION = "The OpenCitations Indexes: http://opencitations.net/"
    DATASET_ID = DAG_ID
    QUEUE = "remote_queue"
    RETRIES = 3
    RELEASES_TOPIC_NAME = "releases"

    TASK_ID_CHECK_DEPENDENCIES = "check_dependencies"
    TASK_ID_LIST_RELEASES = f"list_releases"
    TASK_ID_DOWNLOAD = f"download"
    TASK_ID_UPLOAD_DOWNLOADED = f"upload_downloaded"
    TASK_ID_EXTRACT = f"extract"
    TASK_ID_UPLOAD_EXTRACTED = f"upload_extracted"
    TASK_ID_BQ_LOAD = f"bq_load"
    TASK_ID_CLEANUP = f"cleanup"

    @staticmethod
    def check_dependencies(**kwargs):
        """Check that all variables exist that are required to run the DAG.

        :param kwargs: the context passed from the PythonOperator. See
        https://airflow.apache.org/docs/stable/macros-ref.html
        for a list of the keyword arguments that are passed to this argument.
        """

        vars_valid = check_variables(
            AirflowVars.DATA_PATH,
            AirflowVars.PROJECT_ID,
            AirflowVars.DATA_LOCATION,
            AirflowVars.DOWNLOAD_BUCKET,
            AirflowVars.TRANSFORM_BUCKET,
        )
        if not vars_valid:
            raise AirflowException("Required variables are missing")

    @staticmethod
    def list_releases(**kwargs):
        """Task to lists all Open Citations releases for a given time period.

        :param kwargs: the context passed from the BranchPythonOperator. See
        https://airflow.apache.org/docs/stable/macros-ref.html
        for a list of the keyword arguments that are passed to this argument.
        :return: the identifier of the task to execute next.
        """

        start_date = kwargs["execution_date"]
        end_date = kwargs["next_execution_date"].subtract(microseconds=1)
        releases = list_open_citations_releases(start_date=start_date, end_date=end_date)
        project_id = Variable.get(AirflowVars.PROJECT_ID)

        # Check if we can skip any releases
        releases_out = []
        for release in releases:
            table_id = bigquery_sharded_table_id(OpenCitationsTelescope.DAG_ID, release.release_date)

            if bigquery_table_exists(project_id, OpenCitationsTelescope.DATASET_ID, table_id):
                logging.info(
                    f"Skipping as table exists for {release.release_name} release: "
                    f"{project_id}.{OpenCitationsTelescope.DATASET_ID}.{table_id}"
                )
            else:
                logging.info(
                    f"Table doesn't exist yet, processing Open Citations {release.release_date} "
                    f"release in this workflow"
                )
                releases_out.append(release)

        continue_dag = len(releases_out)
        if continue_dag:
            # Push messages
            ti: TaskInstance = kwargs["ti"]
            ti.xcom_push(OpenCitationsTelescope.RELEASES_TOPIC_NAME, releases_out, start_date)
        return continue_dag

    @staticmethod
    def download(**kwargs):
        """Task to download the Open Citations releases for a given time period.

        :param kwargs: the context passed from the PythonOperator. See
        https://airflow.apache.org/docs/stable/macros-ref.html
        for a list of the keyword arguments that are passed to this argument.
        :return: None.
        """

        # Pull messages
        ti: TaskInstance = kwargs["ti"]
        releases = pull_releases(ti)

        # Download and extract each release posted this month
        for release in releases:
            os.makedirs(release.download_path, exist_ok=True)
            for file in release.files:
                download_file(url=file.download_url, filename=file.name, hash=file.md5hash, hash_algorithm="md5")

    @staticmethod
    def upload_downloaded(**kwargs):
        """Task to upload the downloaded Open Citations releases for a given month.

        :param kwargs: the context passed from the PythonOperator. See
        https://airflow.apache.org/docs/stable/macros-ref.html
        for a list of the keyword arguments that are passed to this argument.
        :return: None.
        """

        # Get bucket name
        bucket_name = Variable.get(AirflowVars.DOWNLOAD_BUCKET)

        # Pull messages
        ti: TaskInstance = kwargs["ti"]
        releases = pull_releases(ti)

        # Upload each release
        file_paths = []
        blob_names = []
        for release in releases:
            for file in release.files:
                file_path = os.path.join(release.download_path, file.name)
                blob_name = file.download_blob_name
                file_paths.append(file_path)
                blob_names.append(blob_name)

        success = upload_files_to_cloud_storage(bucket_name, blob_names, file_paths)
        if not success:
            raise AirflowException("Problem uploading files")

    @staticmethod
    def extract(**kwargs):
        """Task to extract the downloaded Open Citations releases for a given month.

        :param kwargs: the context passed from the PythonOperator. See
        https://airflow.apache.org/docs/stable/macros-ref.html
        for a list of the keyword arguments that are passed to this argument.
        :return: None.
        """

        # Pull messages
        ti: TaskInstance = kwargs["ti"]
        releases = pull_releases(ti)

        for release in releases:
            os.makedirs(release.extract_path, exist_ok=True)

            for file in release.files:
                file_path = os.path.join(release.download_path, file.name)
                cmd = f"unzip {file_path} -d {release.extract_path}"
                logging.info(f"Running command: {cmd}")
                p = subprocess.Popen(
                    cmd, shell=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE, executable="/bin/bash"
                )
                stdout, stderr = wait_for_process(p)

                if stdout:
                    logging.info(stdout)

                if stderr:
                    raise AirflowException(f"bash command failed for {file_path}, {release.release_date}: {stderr}")

                logging.info(f"File extracted to: {release.extract_path}")

    @staticmethod
    def upload_extracted(**kwargs):
        """Task to upload the extracted Open Citations releases for a given month.

        :param kwargs: the context passed from the PythonOperator. See
        https://airflow.apache.org/docs/stable/macros-ref.html
        for a list of the keyword arguments that are passed to this argument.
        :return: None.
        """

        # Get bucket name
        bucket_name = Variable.get(AirflowVars.TRANSFORM_BUCKET)

        # Pull messages
        ti: TaskInstance = kwargs["ti"]
        releases = pull_releases(ti)

        # Upload each release
        file_paths = []
        blob_names = []
        for release in releases:
            release_file_paths = glob.glob(f"{release.extract_path}/*.csv")
            for file_path in release_file_paths:
                file_name = os.path.basename(file_path)
                blob_name = f"telescopes/{OpenCitationsTelescope.DAG_ID}/{release.release_name}/{file_name}"
                file_paths.append(file_path)
                blob_names.append(blob_name)

        logging.info(f"file_paths: {file_paths}")
        logging.info(f"blob_names: {blob_names}")

        success = upload_files_to_cloud_storage(bucket_name, blob_names, file_paths)
        if not success:
            raise AirflowException("Problem uploading files")

    @staticmethod
    def load_to_bq(**kwargs):
        """Load transformed Open Citations release into BigQuery.

        :param kwargs: the context passed from the PythonOperator. See
        https://airflow.apache.org/docs/stable/macros-ref.html
        for a list of the keyword arguments that are passed to this argument.
        :return: None.
        """

        # Pull messages
        ti: TaskInstance = kwargs["ti"]
        releases = pull_releases(ti)

        # Get variables
        project_id = Variable.get(AirflowVars.PROJECT_ID)
        data_location = Variable.get(AirflowVars.DATA_LOCATION)
        bucket_name = Variable.get(AirflowVars.TRANSFORM_BUCKET)

        # Create dataset
        dataset_id = OpenCitationsTelescope.DATASET_ID
        create_bigquery_dataset(project_id, dataset_id, data_location, OpenCitationsTelescope.DESCRIPTION)
        table_name = OpenCitationsTelescope.DAG_ID

        # Load each release
        for release in releases:
            table_id = bigquery_sharded_table_id(table_name, release.release_date)

            # Select schema file based on release date
            analysis_schema_path = schema_folder()
            schema_file_path = find_schema(analysis_schema_path, table_name, release.release_date)
            if schema_file_path is None:
                logging.error(
                    f"No schema found with search parameters: analysis_schema_path={analysis_schema_path}, "
                    f"table_name={table_name}, release_date={release.release_date}"
                )
                exit(os.EX_CONFIG)

            # Load BigQuery table
            uri = f"gs://{bucket_name}/{release.transformed_blob_path}"
            logging.info(f"URI: {uri}")
            success = load_bigquery_table(
                uri,
                dataset_id,
                data_location,
                table_id,
                schema_file_path,
                SourceFormat.CSV,
                csv_field_delimiter=",",
                csv_quote_character='"',
                csv_skip_leading_rows=1,
                csv_allow_quoted_newlines=True,
            )
            if not success:
                raise AirflowException("bq_load task: data failed to load data into BigQuery")

    @staticmethod
    def cleanup(**kwargs):
        """Delete files of downloaded and extracted Open Citations releases.

        :param kwargs: the context passed from the PythonOperator. See
        https://airflow.apache.org/docs/stable/macros-ref.html
        for a list of the keyword arguments that are passed to this argument.
        :return: None.
        """

        # Pull messages
        ti: TaskInstance = kwargs["ti"]
        releases = pull_releases(ti)

        for release in releases:
            logging.info(f"Removing downloaded files for release: {release.release_name}")
            download_path = release.download_path
            try:
                shutil.rmtree(download_path)
            except FileNotFoundError as e:
                logging.warning(f"No such file or directory {download_path}: {e}")

            logging.info(f"Removing extracted files for release: {release.release_name}")
            extract_path = release.extract_path
            try:
                shutil.rmtree(extract_path)
            except FileNotFoundError as e:
                logging.warning(f"No such file or directory {extract_path}: {e}")
