# Copyright 2021 Edward Leardi. All Rights Reserved.
#
# Copyright 2020 IBM Corp. All Rights Reserved.
#
# 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.
#
# This file has been modified by Edward Leardi to support downloading & extracting new dataset formats
#

"Dataset downloading and loading functionality."


from enum import IntFlag
import hashlib
import os
import pathlib
import shutil
from typing import Any, Callable, Dict, Iterable, Optional

import requests

from . import typing as typing_
from .loaders import FormatLoaderMap
from .loaders._format_loader_map import load_data_files
from .schema import SchemaDict
from ._extractors import extract_data_files, verify_data_files
from ._lock import DirectoryLock


class Dataset:
    """Models a particular dataset version along with download & load functionality.

    :param schema: Schema dict of a particular dataset version.
    :param data_dir: Directory to/from which the dataset should be downloaded/loaded from. The path can be either
        absolute or relative to the current working directory, but will be converted to the absolute path immediately
        upon initialization.
    :param mode: Mode with which to treat a dataset. Available options are:
        :attr:`Dataset.InitializationMode.LAZY`, :attr:`Dataset.InitializationMode.DOWNLOAD_ONLY`,
        :attr:`Dataset.InitializationMode.LOAD_ONLY`, and :attr:`Dataset.InitializationMode.DOWNLOAD_AND_LOAD`.
    :raises ValueError: An invalid ``mode`` was specified for handling the dataset.

    Example:

    >>> from tempfile import TemporaryDirectory
    >>> import nourish
    >>> from nourish import schema
    >>> dataset_schema = schema.DatasetSchema('./tests/schemata/datasets.yaml')
    >>> jfk_schema_dict = dataset_schema.export_schema('datasets', 'noaa_jfk', '1.1.4')
    >>> jfk_data_dir = TemporaryDirectory()
    >>> jfk_dataset = Dataset(schema=jfk_schema_dict, data_dir=jfk_data_dir.name)
    >>> jfk_dataset.download()
    >>> data = jfk_dataset.load()
    >>> data['jfk_weather_cleaned'].shape
    (75119, 16)
    >>> jfk_dataset.delete()  # The directory jfk_data_dir is deleted here
    >>> jfk_dataset.is_downloaded()
    False
    """

    # Note: we use a convention for attribute pairs that return the same pathlib.Path where a trailing "_" indicates
    # the path should be used for reading data (e.g. checking the path exists) while no trailing "_" indicates the path
    # is for writing data (e.g. creating the file)

    class InitializationMode(IntFlag):
        """Enum class that acts as `mode` for :class:`Dataset`.
        """
        LAZY = 0
        DOWNLOAD_ONLY = 1
        LOAD_ONLY = 2
        DOWNLOAD_AND_LOAD = 3

    def __init__(self, schema: SchemaDict,
                 data_dir: typing_.PathLike,  *,
                 mode: InitializationMode = InitializationMode.LAZY) -> None:
        """Constructor method.
        """

        self._schema: SchemaDict = schema
        self._data_dir_: pathlib.Path = pathlib.Path(os.path.abspath(data_dir))
        self._data: Optional[Dict[str, Any]] = None
        # Put directory lock under self._nourish_dir. We use self._nourish_dir_ instead of self._nourish_dir because we
        # don't want to have the directory created in lazy mode upon construction of a Dataset object.
        self._lock: DirectoryLock = DirectoryLock(self._nourish_dir_)

        if not isinstance(mode, Dataset.InitializationMode):
            raise ValueError(f'{mode} not a valid mode')

        if mode & Dataset.InitializationMode.DOWNLOAD_ONLY:
            self.download()
        if mode & Dataset.InitializationMode.LOAD_ONLY:
            self.load()

    @property
    def _data_dir(self) -> pathlib.Path:
        "Same as :attr:`_data_dir_`, but create it if it does not exist."
        if not self._data_dir_.exists():
            self._data_dir_.mkdir(parents=True)
        elif not self._data_dir_.is_dir():  # self._data_dir_ exists and is not a directory
            raise NotADirectoryError(f'"{self._data_dir_}" exists and is not a directory.')
        return self._data_dir_

    @property
    def _nourish_dir_(self) -> pathlib.Path:
        "Cache, metainfo, etc. directory used by this class."
        return self._data_dir_ / '.nourish.dataset'

    @property
    def _nourish_dir(self) -> pathlib.Path:
        "Same as :attr:`_nourish_dir_`, but create it if it does not exist."
        if not self._nourish_dir_.exists():
            self._nourish_dir_.mkdir(parents=True)
        elif not self._nourish_dir_.is_dir():  # nourish_dir exists and is not a directory
            raise NotADirectoryError(f'"{self._nourish_dir_}" exists and is not a directory.')
        return self._nourish_dir_

    @property
    def _file_list_file_(self) -> pathlib.Path:
        "Path to the file that stores the list of files in the downloaded dataset."
        return self._nourish_dir_ / 'files.list'

    @property
    def _file_list_file(self) -> pathlib.Path:
        "Same as :attr:`_file_list_file_`, but create the parent directory if it does not exist."
        return self._nourish_dir / 'files.list'

    def download(self,
                 check: bool = True) -> None:
        """Downloads, extracts, and removes dataset archive. It adds a directory write lock during execution.

        :param check: Check to make sure the data files are not already present in :attr:`._data_dir` (passed in via
            ``data_dir`` in the constructor :class:`Dataset`) by running :meth:`.is_downloaded`. If set to ``True``,
            raise an error if they are present and prevent a subsequent download. Set to ``False`` to remove this
            safeguard, and subsequent calls to :meth:`.download` will then overwrite data files if they were previously
            downloaded to :attr:`._data_dir`.
        :raises RuntimeError: The dataset was previously downloaded as indicated by :meth:`.is_downloaded`
            returning ``True``.
        :raises NotADirectoryError: :attr:`Dataset._data_dir` (passed in via ``data_dir`` in the constructor
            :class:`Dataset`) points to an existing file that is not a directory.
        :raises OSError: The SHA512 checksum of a downloaded dataset doesn't match the expected checksum.
        :raises exceptions.DirectoryLockAcquisitionError: Failed to acquire the directory lock.
        """

        if check and self.is_downloaded():
            raise RuntimeError(f'{self.__class__.__name__}.download() was previously called. To overwrite existing '
                               f'data files, rerun {self.__class__.__name__}.download() with ``check`` set to '
                               f'``False``.')

        download_url = self._schema['download_url']
        download_file_name = pathlib.Path(os.path.basename(download_url))

        with self._lock.locking_with_exception(write=True):
            archive_fp = self._nourish_dir / download_file_name
            response = requests.get(download_url, stream=True)
            archive_fp.write_bytes(response.content)

            computed_hash = hashlib.sha512(archive_fp.read_bytes()).hexdigest()
            actual_hash = self._schema['sha512sum']
            if not actual_hash == computed_hash:
                raise OSError(f'{archive_fp} has a SHA512 checksum of: ({computed_hash}) '
                              f'which is different from the expected SHA512 checksum of: ({actual_hash}) '
                              f'the file may by corrupted.')

            extract_data_files(path=archive_fp, data_dir=self._data_dir, file_list_file=self._file_list_file)
            os.remove(archive_fp)

    def load(self,
             subdatasets: Optional[Iterable[str]] = None,
             format_loader_map: Optional[FormatLoaderMap] = None,
             check: bool = True) -> Dict[str, Any]:
        """Load data files to RAM. It adds a directory read lock during execution.

        :param subdatasets: The subdatasets to load. ``None`` means all subdatasets.
        :param format_loader_map: The :class:`.loaders.FormatLoaderMap` object that determines which loader to use.
        :param check: Check to make sure the data files are currently present in :attr:`._data_dir` (passed in via
            ``data_dir`` in the constructor :class:`Dataset`) before loading them by running :meth:`.is_downloaded`.
            If set to ``True``, raise an error if they are missing and prevent attempting to load them. Set to ``False``
            to remove this safeguard.
        :raises RuntimeError: The dataset was not previously downloaded as indicated by :meth:`.is_downloaded`
            returning ``False``.
        :raises FileNotFoundError: The dataset files for a particular subdataset are not found on the disk. Either this
            is because :func:`~Dataset.download` was never called, or the dataset was only partially downloaded.
        :raises exceptions.DirectoryLockAcquisitionError: Failed to acquire the directory lock.
        :return: Loaded data objects. Same as :attr:`.data`.
        """

        if subdatasets is None:
            subdatasets = self._schema['subdatasets'].keys()

        if check and not self.is_downloaded():
            raise RuntimeError(f'Downloaded data files are not present in {self._data_dir_} or are corrupted.')

        with self._lock.locking_with_exception(write=False):
            self._data = {}
            for subdataset in subdatasets:
                subdataset_schema = self._schema['subdatasets'][subdataset]
                try:
                    self._data[subdataset] = load_data_files(fmt=subdataset_schema['format'],
                                                             path=self._data_dir / subdataset_schema['path'],
                                                             format_loader_map=format_loader_map)
                except FileNotFoundError as e:
                    self._data = None
                    raise FileNotFoundError(
                        f'Failed to load subdataset "{subdataset}" because some files are not found. '
                        f'Did you forget to call {self.__class__.__name__}.download()?\nCaused by:\n{e}')

        return self.data

    def delete(self, *, force: bool = False) -> None:
        """Clear the data directory. It adds a directory write lock before deletion during execution if the data
        directory exists.

        :param force: If ``True``, delete the directory even if directory locks are present.
        :raises exceptions.DirectoryLockAcquisitionError: Failed to acquire the directory lock.
        """

        if self._data_dir_.exists():
            if force:
                lock_func: Callable = self._lock.locking
            else:
                lock_func = self._lock.locking_with_exception
            with lock_func(write=True):
                shutil.rmtree(self._data_dir_)

    @property
    def data(self) -> Dict[str, Any]:
        """Access loaded data objects."""
        if self._data is None:
            raise RuntimeError(f'Data has not been downloaded and/or loaded yet. Call '
                               f'{self.__class__.__name__}.download() to download data, call '
                               f'{self.__class__.__name__}.load() to load data.')
        # We don't copy here because it is too expensive and users may actually want to update the datasets and it
        # doesn't cause security issues as in the Schema class
        return self._data

    def is_downloaded(self) -> bool:
        """Check to see if the dataset was downloaded. We determine this by comparing the extracted file tree with the file
        list :meth:`._file_list_file` (their existence, types, and sizes). In this way, if the extraction of the archive
        failed, this should return ``False`` and the user would not be misled. For performance reasons, we do not
        examine the content of the extracted files.

        :return: ``True`` if the dataset has been downloaded and ``False`` otherwise.

        .. warning::

            :meth:`.is_downloaded` will search for the dataset files in :attr:`Dataset._data_dir` (passed in via
            ``data_dir`` in the constructor :class:`Dataset`). If after downloading, you manipulate the data files
            outside the control of this library, this method may produce unexpected behavior.
        """

        # The method to detect whether the dataset has been downloaded can certainly be improved by balancing how much
        # to examine and how much time it takes to do so, considering some heuristics such as the dataset size, the
        # number of files, etc. The method used here should be able to strike a good balance for most cases and should
        # be good enough for the first release.

        if not self._file_list_file_.exists():
            # File not found, may not have finished downloading at all and we treat it as so. We can't control users'
            # own tweaking with the directory.
            return False
        return verify_data_files(data_dir=self._data_dir, file_list_file=self._file_list_file_)
