import copy
import datetime
import functools
import json
import logging
import os
import shutil
import tempfile
import warnings
from urllib.parse import urlparse

import google.api_core.exceptions
import pkg_resources
import requests
from google_crc32c import Checksum

from octue.resources.label import LabelSet
from octue.resources.tag import TagDict


try:
    import h5py
except ModuleNotFoundError:
    pass

from octue.cloud import storage
from octue.cloud.storage import GoogleCloudStorageClient
from octue.exceptions import CloudLocationNotSpecified, ReadOnlyResource
from octue.migrations.cloud_storage import translate_bucket_name_and_path_in_bucket_to_cloud_path
from octue.mixins import Filterable, Hashable, Identifiable, Labelable, Metadata, Serialisable, Taggable
from octue.mixins.hashable import EMPTY_STRING_HASH_VALUE
from octue.utils.decoders import OctueJSONDecoder
from octue.utils.metadata import METADATA_FILENAME, load_local_metadata_file, overwrite_local_metadata_file


logger = logging.getLogger(__name__)


OCTUE_METADATA_NAMESPACE = "octue"


class Datafile(Labelable, Taggable, Serialisable, Identifiable, Hashable, Filterable, Metadata):
    """A representation of a data file on the Octue system. Metadata for the file is obtained from its corresponding
    cloud object or a local `.octue` metadata file, if present. If no stored metadata is available, it can be set during
    or after instantiation.

    :param str|None path: The path of this file locally or in the cloud, which may include folders or subfolders, within the dataset
    :param str|None local_path: If a cloud path is given as the `path` parameter, this is the path to an existing local file that is known to be in sync with the cloud object
    :param str|None cloud_path: If a local path is given for the `path` parameter, this is a cloud path to keep in sync with the local file
    :param datetime.datetime|int|float|None timestamp: A posix timestamp associated with the file, in seconds since epoch, typically when it was created but could relate to a relevant time point for the data
    :param str id: The Universally Unique ID of this file (checked to be valid if not None, generated if None)
    :param dict|octue.resources.tag.TagDict|None tags: key-value pairs with string keys conforming to the Octue tag format (see TagDict)
    :param iter(str)|octue.resources.label.LabelSet|None labels: Space-separated string of labels relevant to this file
    :param str mode: if using as a context manager, open the datafile for reading/editing in this mode (the mode options are the same as for the builtin `open` function)
    :param bool update_metadata: if using as a context manager and this is `True`, update the stored metadata of the datafile when the context is exited
    :param bool hypothetical: if `True`, ignore any metadata stored for this datafile locally or in the cloud and use whatever is given at instantiation
    :return None:
    """

    _METADATA_ATTRIBUTES = ("id", "timestamp", "tags", "labels")

    _SERIALISE_FIELDS = (
        *_METADATA_ATTRIBUTES,
        "name",
        "path",
        "cloud_path",
        "_cloud_metadata",
    )

    def __init__(
        self,
        path,
        local_path=None,
        cloud_path=None,
        timestamp=None,
        id=None,
        tags=None,
        labels=None,
        mode="r",
        update_metadata=True,
        hypothetical=False,
        **kwargs,
    ):
        if "update_cloud_metadata" in kwargs:
            warnings.warn(
                message=(
                    "The `Datafile` instantiation parameter `update_cloud_metadata` has been deprecated and renamed to "
                    "`update_metadata`. The old name will become unavailable soon, so please update to the new one. It "
                    "has been translated for now."
                ),
                category=DeprecationWarning,
            )

            update_metadata = kwargs["update_cloud_metadata"]

        super().__init__(
            id=id,
            name=kwargs.pop("name", None),
            immutable_hash_value=kwargs.pop("immutable_hash_value", None),
            tags=tags,
            labels=labels,
        )

        self.timestamp = timestamp
        self._open_attributes = {"mode": mode, "update_metadata": update_metadata, **kwargs}
        self._local_path = None
        self._cloud_path = None
        self._cloud_metadata = {}

        if storage.path.is_cloud_path(path):
            self._instantiate_from_cloud_object(path, local_path, ignore_stored_metadata=hypothetical)
        else:
            self._instantiate_from_local_path(path, cloud_path, ignore_stored_metadata=hypothetical)

        if hypothetical:
            logger.debug("Ignored stored metadata for %r.", self)
        else:
            if self.metadata(use_octue_namespace=False, include_sdk_version=False) != {
                "id": id or self.id,
                "timestamp": timestamp,
                "tags": TagDict(tags),
                "labels": LabelSet(labels),
            }:
                logger.warning(
                    "Overriding metadata given at instantiation with stored metadata for %r - set `hypothetical` to "
                    "`True` at instantiation to avoid this.",
                    self,
                )

    @classmethod
    def deserialise(cls, serialised_datafile, from_string=False):
        """Deserialise a Datafile from a dictionary or JSON string.

        :param dict|str serialised_datafile:
        :param bool from_string:
        :return Datafile:
        """
        if from_string:
            serialised_datafile = json.loads(serialised_datafile, cls=OctueJSONDecoder)
        else:
            serialised_datafile = copy.deepcopy(serialised_datafile)

        cloud_metadata = serialised_datafile.pop("_cloud_metadata", {})
        datafile = Datafile(**serialised_datafile)
        datafile._cloud_metadata = cloud_metadata
        return datafile

    @property
    def name(self):
        """Get the name of the datafile.

        :return str:
        """
        if self._name:
            return self._name

        if self.exists_in_cloud:
            return str(storage.path.split(self.cloud_path)[-1]).split("?")[0]

        return str(os.path.split(self.local_path)[-1])

    @property
    def extension(self):
        """Get the extension of the datafile.

        :return str:
        """
        return os.path.splitext(self.name)[-1].strip(".")

    @property
    def cloud_path(self):
        """Get the cloud path of the datafile.

        :return str|None:
        """
        return self._cloud_path

    @cloud_path.setter
    def cloud_path(self, path):
        """Set the cloud path of the datafile.

        :param str|None path:
        :return None:
        """
        if path is None:

            if not self.exists_locally:
                raise CloudLocationNotSpecified(
                    "The cloud path cannot be reset because this datafile only exists in the cloud."
                )

            self._cloud_path = None

        else:
            self.to_cloud(cloud_path=path)

    @property
    def cloud_protocol(self):
        """Get the cloud protocol of the datafile if it exists in the cloud (e.g. "gs" for a cloud path of
        "gs://my-bucket/my-file.txt").

        :return str|None:
        """
        if not self.exists_in_cloud:
            return None
        return urlparse(self.cloud_path).scheme

    @property
    def bucket_name(self):
        """Get the name of the bucket the datafile exists in if it exists in the cloud.

        :return str|None:
        """
        if self.cloud_path:
            return storage.path.split_bucket_name_from_cloud_path(self.cloud_path)[0]
        return None

    @property
    def path_in_bucket(self):
        """Get the path of the datafile in its bucket if it exists in the cloud.

        :return str|None:
        """
        if self.cloud_path:
            return storage.path.split_bucket_name_from_cloud_path(self.cloud_path)[1]
        return None

    @property
    def cloud_hash_value(self):
        """Get the hash value of the datafile according to its cloud file.

        :return str|None: `None` if no cloud metadata is available
        """
        return self._cloud_metadata.get("crc32c")

    @property
    def timestamp(self):
        """Get the timestamp of the datafile.

        :return float:
        """
        return self._timestamp

    @timestamp.setter
    def timestamp(self, value):
        """Set the datafile's timestamp.

        :param datetime.datetime|int|float|None value:
        :raise TypeError: if value is of an incorrect type
        :return None:
        """
        if isinstance(value, datetime.datetime) or value is None:
            self._timestamp = value
        elif isinstance(value, (int, float)):
            self._timestamp = datetime.datetime.fromtimestamp(value)
        else:
            raise TypeError(
                f"timestamp should be a datetime.datetime instance, an int, a float, or None; received {value!r}"
            )

    @property
    def posix_timestamp(self):
        """Get the timestamp of the datafile in posix format.

        :return float:
        """
        if self.timestamp is None:
            return None

        return self.timestamp.timestamp()

    @property
    def _last_modified(self):
        """Get the date/time the file was last modified in units of seconds since epoch (posix time).

        :return float|None:
        """
        if self.exists_in_cloud:
            last_modified = self._cloud_metadata.get("updated")

            if last_modified is None:
                return None

            return last_modified.timestamp()

        try:
            return os.path.getmtime(self.local_path)
        except FileNotFoundError:
            return None

    @property
    def size_bytes(self):
        """Get the size of the datafile in bytes.

        :return float|None:
        """
        if self.exists_in_cloud:
            return self._cloud_metadata.get("size")

        try:
            return os.path.getsize(self.local_path)
        except FileNotFoundError:
            return None

    @property
    def exists_in_cloud(self):
        """Return `True` if the file exists in the cloud.

        :return bool:
        """
        return self.cloud_path is not None

    @property
    def exists_locally(self):
        """Return `True` if the file exists locally.

        :return bool:
        """
        return self._local_path is not None

    @property
    def local_path(self):
        """Get the local path for the datafile, downloading it from the cloud to a temporary file if necessary. If
        downloaded, the local path is added to a cache to avoid downloading again in the same runtime.

        :return str:
        """
        if self._local_path:
            return self._local_path

        return self.download()

    @property
    def path(self):
        """Alias to the `local_path` property.

        :return str:
        """
        return self.local_path

    @local_path.setter
    def local_path(self, path):
        """Set the local path of the datafile and:
        - If it exists in the cloud, download the contents of the corresponding cloud file to the new local path
        - If it only exists locally, copy the contents of the old local path to the new local path

        :param str path:
        :raise octue.exceptions.CloudLocationNotSpecified: if `path` is `None` and the datafile doesn't exist in the cloud
        :raise FileExistsError: if the new path corresponds to an existing local file
        :return None:
        """
        if path is None:
            if not self.exists_in_cloud:
                raise CloudLocationNotSpecified(
                    "The local path cannot be reset because this datafile only exists locally."
                )

            self._local_path = None
            return

        if os.path.exists(path):
            raise FileExistsError(
                "Only a path not corresponding to an existing file can be used. This is because the contents of the "
                "existing file would overwrite the existing file at the given path."
            )

        if self.exists_in_cloud:
            GoogleCloudStorageClient().download_to_file(local_path=path, cloud_path=self.cloud_path)
        else:
            os.makedirs(os.path.split(path)[0], exist_ok=True)
            shutil.copy(self._local_path, path)

        self._local_path = os.path.abspath(path)

    @property
    def open(self):
        """Get a context manager for handling the opening and closing of the datafile for reading/editing.

        :return type: the class octue.resources.datafile._DatafileContextManager
        """
        return functools.partial(_DatafileContextManager, self)

    @property
    def _local_metadata_path(self):
        """Get the path to the datafile's local metadata file (if the datafile exists locally).

        :return str|None:
        """
        if not self.exists_locally:
            return None

        return os.path.join(os.path.dirname(self._local_path), METADATA_FILENAME)

    def __enter__(self):
        self._open_context_manager = self.open(**self._open_attributes)
        return self, self._open_context_manager.__enter__()

    def __exit__(self, exc_type, exc_val, exc_tb):
        self._open_context_manager.__exit__(exc_type, exc_val, exc_tb)
        del vars(self)["_open_context_manager"]

    def __lt__(self, other):
        if not isinstance(other, Datafile):
            raise TypeError(f"An object of type {type(self)} cannot be compared with {type(other)}.")
        return self.name < other.name

    def __gt__(self, other):
        if not isinstance(other, Datafile):
            raise TypeError(f"An object of type {type(self)} cannot be compared with {type(other)}.")
        return self.name > other.name

    def to_cloud(self, cloud_path=None, bucket_name=None, path_in_bucket=None, update_cloud_metadata=True):
        """Upload a datafile to Google Cloud Storage.

        :param str|None cloud_path: full path to cloud storage location to store datafile at (e.g. `gs://bucket_name/path/to/file.csv`)
        :param bool update_cloud_metadata: if `True`, update the metadata of the datafile in the cloud at upload time
        :return str: gs:// path for datafile
        """
        if bucket_name:
            cloud_path = translate_bucket_name_and_path_in_bucket_to_cloud_path(bucket_name, path_in_bucket)

        cloud_path = self._get_cloud_location(cloud_path)

        self._get_cloud_metadata()

        storage_client = GoogleCloudStorageClient()

        # If the there is no cloud file or if the datafile's file has been changed locally, overwrite its cloud copy.
        if not storage_client.exists(cloud_path) or self.cloud_hash_value != self.hash_value:
            storage_client.upload_file(
                local_path=self.local_path,
                cloud_path=cloud_path,
                metadata=self.metadata(),
            )

        if update_cloud_metadata:
            # If the datafile's metadata has been changed locally, update the cloud file's metadata.
            local_metadata = self.metadata(use_octue_namespace=False)

            if self._cloud_metadata.get("custom_metadata") != local_metadata:
                self.update_cloud_metadata()

        return self.cloud_path

    def download(self, local_path=None):
        """Download the file from the cloud to the given local path or a random temporary path.

        :param str|None local_path:
        :raise CloudLocationNotSpecified: if the datafile does not exist in the cloud
        :return str: path to local file
        """
        if not self.exists_in_cloud:
            raise CloudLocationNotSpecified("Cannot download a file that doesn't exist in the cloud.")

        # Avoid downloading to a local path if the datafile has already been downloaded to it.
        if (local_path is None and self._local_path is not None) or (
            local_path is not None and local_path == self._local_path
        ):
            return self._local_path

        if local_path is not None:
            self._local_path = os.path.abspath(local_path)
        else:
            self._local_path = tempfile.NamedTemporaryFile(delete=False).name

        # Download from a URL.
        if storage.path.is_url(self.cloud_path):
            with open(self._local_path, "wb") as f:
                f.write(requests.get(self.cloud_path).content)

        # Download from a cloud URI.
        else:
            try:
                GoogleCloudStorageClient().download_to_file(local_path=self._local_path, cloud_path=self.cloud_path)

            except google.api_core.exceptions.NotFound as e:
                # If in reading mode, raise an error if no file exists at the path; if in a writing mode, create a new file.
                if self._open_attributes["mode"] == "r":
                    raise e

        # Now use hash value of local file instead of cloud file.
        self.reset_hash()
        return self._local_path

    def metadata(self, include_sdk_version=True, use_octue_namespace=True):
        """Get the datafile's metadata in a serialised form (i.e. the attributes `id`, `timestamp`, `labels`, `tags`,
        and `sdk_version`).

        :param bool include_sdk_version: if `True`, include the `octue` version that instantiated the datafile in the metadata
        :param bool use_octue_namespace: if `True`, prefix metadata names with "octue__"
        :return dict:
        """
        metadata = super().metadata(include_sdk_version=include_sdk_version)

        if not use_octue_namespace:
            return metadata

        return {f"{OCTUE_METADATA_NAMESPACE}__{key}": value for key, value in metadata.items()}

    def update_metadata(self):
        """If the datafile is cloud-based, update its cloud metadata; otherwise, update its local metadata.

        :return None:
        """
        if self.exists_in_cloud:
            self.update_cloud_metadata()
            return

        self.update_local_metadata()

    def update_cloud_metadata(self):
        """Update the cloud metadata for the datafile.

        :return None:
        """
        if not self.cloud_path:
            self._raise_cloud_location_error()

        GoogleCloudStorageClient().overwrite_custom_metadata(metadata=self.metadata(), cloud_path=self.cloud_path)

    def update_local_metadata(self):
        """Create or update the local octue metadata file with the datafile's metadata.

        :return None:
        """
        existing_metadata_records = load_local_metadata_file(self._local_metadata_path)

        if not existing_metadata_records.get("datafiles"):
            existing_metadata_records["datafiles"] = {}

        existing_metadata_records["datafiles"][self.name] = self.metadata(use_octue_namespace=False)
        overwrite_local_metadata_file(data=existing_metadata_records, path=self._local_metadata_path)

    def generate_signed_url(self, expiration=datetime.timedelta(days=7)):
        """Generate a signed URL for the datafile.

        :param datetime.datetime|datetime.timedelta expiration: the amount of time or date after which the URL should expire
        :return str: the signed URL for the datafile
        """
        if not self.exists_in_cloud:
            raise CloudLocationNotSpecified(
                f"{self!r} must exist in the cloud for a signed URL to be generated for it."
            )

        return GoogleCloudStorageClient().generate_signed_url(cloud_path=self.cloud_path, expiration=expiration)

    def _instantiate_from_cloud_object(self, path, local_path, ignore_stored_metadata):
        """Instantiate the datafile from a cloud object.

        :param str path: the cloud path to a cloud object
        :param str|None local_path: a local path to use for opening or modifying the cloud object locally
        :param bool ignore_stored_metadata: if `True`, don't use any metadata stored for this datafile in the cloud
        :return None:
        """
        self._cloud_path = path

        if not ignore_stored_metadata:
            self._use_cloud_metadata()

        if local_path:
            # If there is no file at the given local path or the file is different to the one in the cloud, download
            # the cloud file locally.
            if not os.path.exists(local_path) or self.cloud_hash_value != calculate_hash(local_path):
                self.download(local_path)
            else:
                self._local_path = local_path

    def _instantiate_from_local_path(self, path, cloud_path, ignore_stored_metadata):
        """Instantiate the datafile from a local path.

        :param str path: the path to a local file
        :param str|None cloud_path: a cloud path to upload the datafile to and mirror any changes made to it locally
        :param bool ignore_stored_metadata: if `True`, don't use any metadata stored for this datafile locally
        :return None:
        """
        self._local_path = os.path.abspath(path)

        if not ignore_stored_metadata:
            self._use_local_metadata()

        if cloud_path:
            self.cloud_path = cloud_path

    def _get_cloud_metadata(self):
        """Get the cloud metadata for the datafile and store it without updating the datafile's metadata.

        :return None:
        """
        if not self.cloud_path:
            self._raise_cloud_location_error()

        # Skip getting metadata for now if the cloud path is a signed URL.
        if storage.path.is_url(self.cloud_path):
            return

        cloud_metadata = GoogleCloudStorageClient().get_metadata(cloud_path=self.cloud_path)

        if cloud_metadata:
            cloud_metadata["custom_metadata"] = {
                key.replace(f"{OCTUE_METADATA_NAMESPACE}__", ""): value
                for key, value in cloud_metadata["custom_metadata"].items()
            }

            self._cloud_metadata = cloud_metadata

    def _use_cloud_metadata(self):
        """Update the datafile's metadata from the metadata of the cloud object located at its path. If no metadata is
        stored for the datafile, do nothing.

        :return None:
        """
        self._get_cloud_metadata()
        cloud_custom_metadata = self._cloud_metadata.get("custom_metadata", {})

        if not cloud_custom_metadata:
            return

        self.immutable_hash_value = self.cloud_hash_value or EMPTY_STRING_HASH_VALUE
        self._set_metadata(cloud_custom_metadata)

    def _use_local_metadata(self):
        """Update the datafile's metadata from the local metadata records file. If no metadata is stored for the
        datafile, do nothing.

        :return None:
        """
        existing_metadata_records = load_local_metadata_file(self._local_metadata_path)
        datafile_metadata = existing_metadata_records.get("datafiles", {}).get(self.name, {})

        if not datafile_metadata:
            return

        self._set_metadata(datafile_metadata)

    def _set_metadata(self, metadata):
        """Set the datafile's metadata.

        :param dict metadata:
        :return None:
        """
        for attribute in self._METADATA_ATTRIBUTES:
            if attribute not in metadata:
                continue

            if attribute == "id":
                self._set_id(metadata["id"])
                continue

            setattr(self, attribute, metadata[attribute])

    def _calculate_hash(self):
        """Get the hash of the datafile according to the first of the following methods that is applicable:

        1. The hash of the file at its local path
        2. If it doesn't have a local path, the hash of the file at its cloud path
        3. If it doesn't have either of these, use the empty string hash value

        :return str:
        """
        if self._local_path and os.path.exists(self._local_path):
            # Calculate the hash of the file itself and then pass it to `Hashable` to include the hashes of any
            # attributes named in `self._ATTRIBUTES_TO_HASH`.
            hash_value = calculate_hash(self._local_path)
            return super()._calculate_hash(hash_value)
        else:
            return self.cloud_hash_value or EMPTY_STRING_HASH_VALUE

    def _get_cloud_location(self, cloud_path=None):
        """Get the cloud location details for the bucket, allowing the keyword arguments to override any stored values.
        Once the cloud location details have been determined, update the stored cloud location details.

        :param str|None cloud_path:
        :raise octue.exceptions.CloudLocationNotSpecified: if an exact cloud location isn't provided and isn't available implicitly (i.e. the Datafile wasn't loaded from the cloud previously)
        :return (str, str): project name and cloud path
        """
        cloud_path = cloud_path or self.cloud_path

        if not cloud_path:
            self._raise_cloud_location_error()

        self._cloud_path = cloud_path
        return cloud_path

    def _raise_cloud_location_error(self):
        """Raise an error indicating that the cloud location of the datafile has not yet been specified.

        :raise CloudLocationNotSpecified:
        :return None:
        """
        raise CloudLocationNotSpecified(
            f"{self!r} wasn't previously loaded from the cloud so doesn't have an implicit cloud location - please "
            f"specify its exact location (its project name and cloud path)."
        )


class _DatafileContextManager:
    """A context manager for opening datafiles for reading and writing locally or from the cloud. Its usage is analogous
    to the builtin open context manager. If opening a local datafile in write mode, the manager will attempt to
    determine if the folder path exists and, if not, will create the folder structure required to write the file.

    Usage:
    ```
    my_datafile = Datafile(path='subfolder/subsubfolder/my_datafile.json)
    with my_datafile.open('w') as fp:
        fp.write("{}")
    ```

    This is equivalent to the standard python:
    ```
    my_datafile = Datafile(path='subfolder/subsubfolder/my_datafile.json)
    os.makedirs(os.path.split(my_datafile.local_path)[0], exist_ok=True)
    with open(my_datafile.local_path, 'w') as fp:
        fp.write("{}")
    ```

    :param octue.resources.datafile.Datafile datafile:
    :param str mode: open the datafile for reading/editing in this mode (the mode options are the same as for the builtin `open` function)
    :param bool update_metadata: if this is `True`, update the stored metadata of the datafile when the context is exited
    :return None:
    """

    MODIFICATION_MODES = {"w", "a", "x", "+", "U"}

    def __init__(self, datafile, mode="r", update_metadata=True, **kwargs):
        self.datafile = datafile
        self.mode = mode
        self._update_metadata = update_metadata
        self.kwargs = kwargs
        self._fp = None

    def __enter__(self):
        """Open the datafile, first downloading it from the cloud if necessary.

        :return io.TextIOWrapper:
        """
        if "w" in self.mode:
            if self.datafile.exists_in_cloud and storage.path.is_url(self.datafile.cloud_path):
                raise ReadOnlyResource(f"{self.datafile} is read-only. Change the open mode to 'r' to continue.")

            os.makedirs(os.path.split(self.datafile.local_path)[0], exist_ok=True)

        if self.datafile.extension == "hdf5":
            try:
                pkg_resources.get_distribution("h5py")
                self._fp = h5py.File(self.datafile.local_path, self.mode, **self.kwargs)
            except pkg_resources.DistributionNotFound:
                raise ImportError(
                    "To use datafiles with HDF5 files, please install octue with the 'hdf5' option i.e. "
                    "`pip install octue[hdf5]`."
                )
        else:
            self._fp = open(self.datafile.local_path, self.mode, **self.kwargs)

        return self._fp

    def __exit__(self, *args):
        """Close the datafile, updating the corresponding file in the cloud if necessary and its metadata if
        `self._update_cloud_metadata` is True.

        :return None:
        """
        if self._fp is not None:
            self._fp.close()

        if any(character in self.mode for character in self.MODIFICATION_MODES):

            if self.datafile.exists_in_cloud:
                self.datafile.to_cloud(update_cloud_metadata=self._update_metadata)

            elif self._update_metadata:
                self.datafile.update_local_metadata()


def calculate_hash(path):
    """Calculate the hash of the file at the given path."""
    hash = Checksum()

    with open(path, "rb") as f:
        # Read and update hash value in blocks of 4K.
        for byte_block in iter(lambda: f.read(4096), b""):
            hash.update(byte_block)

    return hash
