# Copyright 2020 Codethink Ltd
# Distributed under the terms of the Modified BSD License.

import abc
import base64
import calendar
import contextlib
import datetime
import hashlib
import logging
import os
import os.path as path
import shutil
import tempfile
import warnings
from contextlib import contextmanager
from os import PathLike
from threading import Lock
from typing import IO, List, Tuple, Union

import aiofiles
import aioshutil
import fsspec
from tenacity import retry
from tenacity.retry import retry_if_exception_type
from tenacity.stop import stop_after_attempt

try:
    import xattr

    has_xattr = True
except ImportError:
    has_xattr = False
    warnings.warn("xattr not installed, not using it for file metadata")

from quetz.errors import ConfigError

File = IO[bytes]

StrPath = Union[str, PathLike]

logger = logging.getLogger("quetz")


class PackageStore(abc.ABC):
    def __init__(self):
        self._download_locks = {}

    def get_download_lock(self, channel: str, destination: str):
        return self._download_locks.get((channel, destination))

    def create_download_lock(self, channel: str, destination: str):
        lock = self._download_locks[(channel, destination)] = Lock()
        return lock

    def delete_download_lock(self, channel: str, destination: str):
        del self._download_locks[(channel, destination)]

    @property
    def kind(self):
        return type(self).__name__

    @property
    @abc.abstractmethod
    def support_redirect(self) -> bool:
        pass

    @abc.abstractmethod
    def create_channel(self, name):
        pass

    @abc.abstractmethod
    def remove_channel(self, name):
        """remove channel recursively"""
        pass

    @abc.abstractmethod
    def list_files(self, channel: str) -> List[str]:
        pass

    @abc.abstractmethod
    def url(self, channel: str, src: str, expires: int = 0) -> str:
        pass

    @abc.abstractmethod
    def add_package_async(self, package: File, channel: str, destination: str):
        pass

    @abc.abstractmethod
    def add_package(self, package: File, channel: str, destination: str):
        pass

    @abc.abstractmethod
    def add_file(
        self, data: Union[str, bytes], channel: str, destination: StrPath
    ) -> None:
        pass

    @abc.abstractmethod
    def serve_path(self, channel, src):
        pass

    @abc.abstractmethod
    def delete_file(self, channel: str, destination: str):
        """remove file from package store"""

    @abc.abstractmethod
    def move_file(self, channel: str, source: str, destination: str):
        """move file from source to destination in package store"""

    @abc.abstractmethod
    def file_exists(self, channel: str, destination: str):
        """Return True if the file exists"""

    @abc.abstractmethod
    def get_filemetadata(self, channel: str, src: str) -> Tuple[int, int, str]:
        """get file metadata: returns (file size, last modified time, etag)"""

    @abc.abstractclassmethod
    def cleanup_temp_files(self, channel: str, dry_run: bool = False):
        """clean up temporary `*.json{HASH}.[bz2|gz]` files from pkgstore"""


# generate a secret token for use with nginx secure link
# similar to https://stackoverflow.com/a/52764346 (thanks @flix on stackoverflow)
def nginx_secure_link(url: str, secret: str, expires=3600):
    future = datetime.datetime.utcnow() + datetime.timedelta(seconds=expires)
    expire_ts = calendar.timegm(future.timetuple())

    hash_string = f"{expire_ts}{url} {secret}".encode('utf-8')
    m = hashlib.md5(hash_string).digest()

    base64_hash = base64.urlsafe_b64encode(m)
    base64_s = base64_hash.decode('utf-8').rstrip('=')
    return base64_s, expire_ts


class LocalStore(PackageStore):
    def __init__(self, config):
        self.fs: fsspec.AbstractFileSystem = fsspec.filesystem("file")
        self.channels_dir = config['channels_dir']
        self.redirect_enabled = config['redirect_enabled']
        self.redirect_endpoint = config['redirect_endpoint']
        self.redirect_secret = config.get('redirect_secret')
        self.redirect_expiration = config.get('redirect_expiration')

        super().__init__()

    @property
    def support_redirect(self):
        return self.redirect_enabled

    @contextmanager
    def _atomic_open(self, channel: str, destination: StrPath, mode="wb") -> IO:
        full_path = path.join(self.channels_dir, channel, destination)
        self.fs.makedirs(path.dirname(full_path), exist_ok=True)

        # Creates a tempfile in the same directory as the target filename.
        # Renames it into place when it's closed.
        fh, tmpname = tempfile.mkstemp(dir=os.path.dirname(full_path), prefix=".")
        f = open(fh, mode)
        try:
            yield f
        except:  # noqa
            f.close()
            os.remove(tmpname)
            raise
        else:
            f.flush()  # Belt and braces (network file systems)
            f.close()
            os.rename(tmpname, full_path)

    def create_channel(self, name):
        self.fs.makedirs(path.join(self.channels_dir, name), exist_ok=True)

    def remove_channel(self, name):
        channel_path = path.join(self.channels_dir, name)
        self.fs.rm(channel_path, recursive=True)

    def add_package(self, package: File, channel: str, destination: str) -> None:
        with self._atomic_open(channel, destination) as f:
            shutil.copyfileobj(package, f)

    async def add_package_async(
        self, package: File, channel: str, destination: str
    ) -> None:
        full_path = path.join(self.channels_dir, channel, destination)
        self.fs.makedirs(path.dirname(full_path), exist_ok=True)

        async with aiofiles.open(full_path, 'wb') as f:
            await f.write(package.read())

    def add_file(
        self, data: Union[str, bytes], channel: str, destination: StrPath
    ) -> None:
        mode = "w" if isinstance(data, str) else "wb"
        with self._atomic_open(channel, destination, mode) as f:
            f.write(data)

    def delete_file(self, channel: str, destination: str):
        self.fs.delete(path.join(self.channels_dir, channel, destination))

    def move_file(self, channel: str, source: str, destination: str):
        self.fs.move(
            path.join(self.channels_dir, channel, source),
            path.join(self.channels_dir, channel, destination),
        )

    def file_exists(self, channel: str, destination: str):
        return self.fs.exists(path.join(self.channels_dir, channel, destination))

    def serve_path(self, channel, src):
        return self.fs.open(path.join(self.channels_dir, channel, src))

    def list_files(self, channel: str):
        channel_dir = os.path.join(self.channels_dir, channel)
        return [os.path.relpath(f, channel_dir) for f in self.fs.find(channel_dir)]

    def url(self, channel: str, src: str, expires=0):
        if self.redirect_enabled:
            # generate url + secret if necessary
            if self.redirect_secret:
                url = path.join(self.channels_dir, channel, src)
                md5hash, expires_by = nginx_secure_link(
                    url, self.redirect_secret, self.redirect_expiration
                )
                return (
                    path.join(self.redirect_endpoint, url)
                    + f"?md5={md5hash}&expires={expires_by}"
                )

            return path.join(self.redirect_endpoint, self.channels_dir, channel, src)
        else:
            return path.abspath(path.join(self.channels_dir, channel, src))

    def get_filemetadata(self, channel: str, src: str):
        filepath = path.abspath(path.join(self.channels_dir, channel, src))
        if not path.exists(filepath):
            raise FileNotFoundError()

        stat_res = os.stat(filepath)
        mtime = stat_res.st_mtime
        msize = stat_res.st_size

        xattr_failed = False
        etag = ""
        if has_xattr:
            try:
                # xattr will fail here if executing on e.g. the tmp filesystem
                attrs = xattr.xattr(filepath)
                attrs['user.testifpermissionok'] = b''
                try:
                    etag = attrs['user.etag'].decode('ascii')
                except KeyError:
                    # calculate md5 sum
                    with self.fs.open(filepath, 'rb') as f:
                        etag = hashlib.md5(f.read()).hexdigest()
                    attrs['user.etag'] = etag.encode('ascii')
            except OSError:
                xattr_failed = True

        elif not has_xattr or xattr_failed:
            etag_base = str(mtime) + "-" + str(msize)
            etag = hashlib.md5(etag_base.encode()).hexdigest()
        else:
            raise RuntimeError

        return (msize, mtime, etag)

    def cleanup_temp_files(self, channel: str, dry_run: bool = False):
        temp_files = []
        for each_end in [".bz2", ".gz"]:
            temp_files.extend(
                self.fs.glob(
                    f"{path.join(self.channels_dir, channel)}/**/*.json?*{each_end}"
                )
            )

        for each_temp_file in temp_files:
            logger.info(f"removing {each_temp_file} from pkgstore")
            if not dry_run:
                self.fs.delete(each_temp_file)


class S3Store(PackageStore):
    def __init__(self, config):
        try:
            import s3fs
        except ModuleNotFoundError:
            raise ModuleNotFoundError("S3 package store requires s3fs module")

        client_kwargs = {}
        url = config.get('url')
        region = config.get("region")
        if url:
            client_kwargs['endpoint_url'] = url
        if region:
            client_kwargs["region_name"] = region

        # When using IAM, key and secret will be empty, so need to pass None
        # to the s3fs constructor
        key = config['key'] if config['key'] != '' else None
        secret = config['secret'] if config['secret'] != '' else None
        self._key = key
        self._secret = secret
        self._client_kwargs = client_kwargs
        self.fs = s3fs.S3FileSystem(
            key=key,
            secret=secret,
            client_kwargs=client_kwargs,
        )

        self.bucket_prefix = config['bucket_prefix']
        self.bucket_suffix = config['bucket_suffix']
        super().__init__()

    @property
    def support_redirect(self):
        return True

    @contextlib.contextmanager
    def _get_fs(self):
        try:
            yield self.fs
        except PermissionError as e:
            raise ConfigError(f"{e} - check configured S3 credentials")

    def _bucket_map(self, name):
        return f"{self.bucket_prefix}{name}{self.bucket_suffix}"

    def create_channel(self, name):
        """Create the bucket if one doesn't already exist

        Parameters
        ----------
        name : str
            The name of the bucket to create on s3
        """
        with self._get_fs() as fs:
            try:
                fs.mkdir(self._bucket_map(name), acl="private")
            except FileExistsError:
                pass

    # we need to retry due to eventual (vs strong) consistency on Openstack Swift
    @retry(stop=stop_after_attempt(5), retry=retry_if_exception_type(OSError))
    def remove_channel(self, name):
        channel_path = self._bucket_map(name)
        self.fs.rm(channel_path, recursive=True, acl="private")

    def add_package(self, package: File, channel: str, destination: str) -> None:
        with self._get_fs() as fs:
            bucket = self._bucket_map(channel)
            with fs.open(path.join(bucket, destination), "wb", acl="private") as pkg:
                # use a chunk size of 10 Megabytes
                shutil.copyfileobj(package, pkg, 10 * 1024 * 1024)

    async def add_package_async(
        self, package: File, channel: str, destination: str
    ) -> None:
        import s3fs

        fs = s3fs.S3FileSystem(
            key=self._key,
            secret=self._secret,
            client_kwargs=self._client_kwargs,
        )
        bucket = self._bucket_map(channel)
        with fs.open(path.join(bucket, destination), "wb", acl="private") as pkg:
            # use a chunk size of 10 Megabytes
            await aioshutil.copyfileobj(package, pkg, 10 * 1024 * 1024)

    def add_file(
        self, data: Union[str, bytes], channel: str, destination: StrPath
    ) -> None:
        if type(data) is str:
            mode = "w"
        else:
            mode = "wb"

        with self._get_fs() as fs:
            bucket = self._bucket_map(channel)
            with fs.open(path.join(bucket, destination), mode, acl="private") as f:
                f.write(data)

    def serve_path(self, channel, src):
        with self._get_fs() as fs:
            return fs.open(path.join(self._bucket_map(channel), src))

    def delete_file(self, channel: str, destination: str):
        channel_bucket = self._bucket_map(channel)

        with self._get_fs() as fs:
            fs.delete(path.join(channel_bucket, destination))

    def move_file(self, channel: str, source: str, destination: str):
        channel_bucket = self._bucket_map(channel)

        with self._get_fs() as fs:
            fs.move(
                path.join(channel_bucket, source),
                path.join(channel_bucket, destination),
            )

    def file_exists(self, channel: str, destination: str):
        channel_bucket = self._bucket_map(channel)
        with self._get_fs() as fs:
            return fs.exists(path.join(channel_bucket, destination))

    def list_files(self, channel: str):
        def remove_prefix(text, prefix):
            if text.startswith(prefix):
                return text[len(prefix) :].lstrip("/")  # noqa: E203
            return text

        channel_bucket = self._bucket_map(channel)

        with self._get_fs() as fs:
            return [remove_prefix(f, channel_bucket) for f in fs.find(channel_bucket)]

    def url(self, channel: str, src: str, expires=3600):
        # expires is in seconds, so the default is 60 minutes!
        with self._get_fs() as fs:
            return fs.url(path.join(self._bucket_map(channel), src), expires)

    def get_filemetadata(self, channel: str, src: str):
        with self._get_fs() as fs:
            filepath = path.join(self._bucket_map(channel), src)
            infodata = fs.info(filepath)

            mtime = infodata['LastModified'].timestamp()
            msize = infodata['size']
            etag = infodata['ETag']

            return (msize, mtime, etag)

    def cleanup_temp_files(self, channel: str, dry_run: bool = False):
        with self._get_fs() as fs:
            temp_files = []
            for each_end in [".bz2", ".gz"]:
                temp_files.extend(
                    fs.glob(f"{self._bucket_map(channel)}/**/*.json?*{each_end}")
                )

            for each_temp_file in temp_files:
                logger.info(f"removing {each_temp_file} from pkgstore")
                if not dry_run:
                    fs.delete(each_temp_file)


class AzureBlobStore(PackageStore):
    def __init__(self, config):
        try:
            import adlfs
        except ModuleNotFoundError:
            raise ModuleNotFoundError("Azure Blob package store requires adlfs module")

        self.storage_account_name = config.get('account_name')
        self.access_key = config.get("account_access_key")
        self.conn_string = config.get("conn_str")

        self.fs = adlfs.AzureBlobFileSystem(
            account_name=self.storage_account_name,
            connection_string=self.conn_string,
            account_key=self.access_key,
        )

        self.container_prefix = config['container_prefix']
        self.container_suffix = config['container_suffix']
        super().__init__()

    @property
    def support_redirect(self):
        return True

    @contextlib.contextmanager
    def _get_fs(self):
        try:
            yield self.fs
        except PermissionError as e:
            raise ConfigError(f"{e} - check configured Azure Blob credentials")

    def _container_map(self, name):
        return f"{self.container_prefix}{name}{self.container_suffix}"

    def create_channel(self, name):
        """Create the container if one doesn't already exist

        Parameters
        ----------
        name : str
            The name of the container to create on azure blob
        """
        with self._get_fs() as fs:
            try:
                fs.mkdir(self._container_map(name))
            except FileExistsError:
                pass

    def remove_channel(self, name):
        channel_path = self._container_map(name)
        with self._get_fs() as fs:
            fs.rm(channel_path, recursive=True)

    def add_package(self, package: File, channel: str, destination: str) -> None:
        with self._get_fs() as fs:
            container = self._container_map(channel)
            with fs.open(path.join(container, destination), "wb") as pkg:
                # use a chunk size of 10 Megabytes
                shutil.copyfileobj(package, pkg, 10 * 1024 * 1024)

    async def add_package_async(
        self, package: File, channel: str, destination: str
    ) -> None:
        import adlfs

        fs = adlfs.AzureBlobFileSystem(
            account_name=self.storage_account_name,
            connection_string=self.conn_string,
            account_key=self.access_key,
        )
        container = self._container_map(channel)
        with fs.open(path.join(container, destination), "wb") as pkg:
            # use a chunk size of 10 Megabytes
            await aioshutil.copyfileobj(package, pkg, 10 * 1024 * 1024)

    def add_file(
        self, data: Union[str, bytes], channel: str, destination: StrPath
    ) -> None:
        if type(data) is str:
            mode = "w"
        else:
            mode = "wb"

        with self._get_fs() as fs:
            container = self._container_map(channel)
            with fs.open(path.join(container, destination), mode) as f:
                f.write(data)

    def serve_path(self, channel, src):
        with self._get_fs() as fs:
            return fs.open(path.join(self._container_map(channel), src))

    def delete_file(self, channel: str, destination: str):
        channel_container = self._container_map(channel)

        with self._get_fs() as fs:
            fs.delete(path.join(channel_container, destination))

    def move_file(self, channel: str, source: str, destination: str):
        channel_container = self._container_map(channel)

        with self._get_fs() as fs:
            fs.move(
                path.join(channel_container, source),
                path.join(channel_container, destination),
            )

    def file_exists(self, channel: str, destination: str):
        channel_container = self._container_map(channel)
        with self._get_fs() as fs:
            return fs.exists(path.join(channel_container, destination))

    def list_files(self, channel: str):
        def remove_prefix(text, prefix):
            if text.startswith(prefix):
                return text[len(prefix) :].lstrip("/")  # noqa: E203
            return text

        channel_container = self._container_map(channel)

        with self._get_fs() as fs:
            return [
                remove_prefix(f, channel_container) for f in fs.find(channel_container)
            ]

    def url(self, channel: str, src: str, expires=3600):
        # expires is in seconds, so the default is 60 minutes!
        with self._get_fs() as fs:
            return fs.url(path.join(self._container_map(channel), src), expires)

    def get_filemetadata(self, channel: str, src: str):
        with self._get_fs() as fs:
            filepath = path.join(self._container_map(channel), src)
            infodata = fs.info(filepath)

            mtime = infodata['last_modified'].timestamp()
            msize = infodata['size']
            etag = infodata['etag']

            return (msize, mtime, etag)

    def cleanup_temp_files(self, channel: str, dry_run: bool = False):
        with self._get_fs() as fs:
            temp_files = []
            for each_end in [".bz2", ".gz"]:
                temp_files.extend(
                    fs.glob(f"{self._container_map(channel)}/**/*.json?*{each_end}")
                )

            for each_temp_file in temp_files:
                logger.info(f"removing {each_temp_file} from pkgstore")
                if not dry_run:
                    fs.delete(each_temp_file)


class GoogleCloudStorageStore(PackageStore):
    def __init__(self, config):
        try:
            import gcsfs
        except ModuleNotFoundError:
            raise ModuleNotFoundError(
                "Google Cloud Storage package store requires gcsfs module"
            )

        self.project = config.get("project")
        self.token = config.get("token")
        self.cache_timeout = config.get("cache_timeout")
        self.region = config.get("region")

        self.fs = gcsfs.GCSFileSystem(
            project=self.project,
            token=self.token if self.token else None,
            cache_timeout=self.cache_timeout,
            default_location=self.region,
        )

        self.bucket_prefix = config['bucket_prefix']
        self.bucket_suffix = config['bucket_suffix']
        super().__init__()

    @property
    def support_redirect(self):
        # `gcsfs` currently doesnt support signing yet. Once this is implemented we
        # can enable this again.
        return True

    @contextlib.contextmanager
    def _get_fs(self):
        try:
            yield self.fs
        except PermissionError as e:
            raise ConfigError(
                f"{e} - check configured Google Cloud Storage credentials"
            )

    def _bucket_map(self, name):
        return f"{self.bucket_prefix}{name}{self.bucket_suffix}"

    def create_channel(self, name):
        """Create the container if one doesn't already exist

        Parameters
        ----------
        name : str
            The name of the container to create on azure blob
        """
        with self._get_fs() as fs:
            bucket_name = self._bucket_map(name)
            if f"{bucket_name}/" not in fs.buckets:
                fs.mkdir(bucket_name)
                fs.invalidate_cache()

    def remove_channel(self, name):
        channel_path = self._bucket_map(name)
        with self._get_fs() as fs:
            fs.rm(channel_path, recursive=True)

    def add_package(self, package: File, channel: str, destination: str) -> None:
        with self._get_fs() as fs:
            container = self._bucket_map(channel)
            with fs.open(path.join(container, destination), "wb") as pkg:
                # use a chunk size of 10 Megabytes
                shutil.copyfileobj(package, pkg, 10 * 1024 * 1024)

    async def add_package_async(
        self, package: File, channel: str, destination: str
    ) -> None:
        import gcsfs

        fs = gcsfs.GCSFileSystem(
            project=self.project,
            token=self.token if self.token else None,
            cache_timeout=self.cache_timeout,
            default_location=self.region,
        )
        container = self._bucket_map(channel)
        with fs.open(path.join(container, destination), "wb") as pkg:
            # use a chunk size of 10 Megabytes
            await aioshutil.copyfileobj(package, pkg, 10 * 1024 * 1024)

    def add_file(
        self, data: Union[str, bytes], channel: str, destination: StrPath
    ) -> None:
        if type(data) is str:
            mode = "w"
        else:
            mode = "wb"

        with self._get_fs() as fs:
            container = self._bucket_map(channel)
            with fs.open(path.join(container, destination), mode) as f:
                f.write(data)

    def serve_path(self, channel, src):
        with self._get_fs() as fs:
            file = fs.open(path.join(self._bucket_map(channel), src))
            info = file.info()
            if info["type"] != "file":
                raise FileNotFoundError()
            return file

    def delete_file(self, channel: str, destination: str):
        channel_container = self._bucket_map(channel)

        with self._get_fs() as fs:
            fs.delete(path.join(channel_container, destination))

    def move_file(self, channel: str, source: str, destination: str):
        channel_container = self._bucket_map(channel)

        with self._get_fs() as fs:
            fs.move(
                path.join(channel_container, source),
                path.join(channel_container, destination),
            )

    def file_exists(self, channel: str, destination: str):
        channel_container = self._bucket_map(channel)
        with self._get_fs() as fs:
            return fs.exists(path.join(channel_container, destination))

    def list_files(self, channel: str):
        def remove_prefix(text, prefix):
            if text.startswith(prefix):
                return text[len(prefix) :].lstrip("/")  # noqa: E203
            return text

        channel_container = self._bucket_map(channel)

        with self._get_fs() as fs:
            return [
                remove_prefix(f, channel_container) for f in fs.find(channel_container)
            ]

    def url(self, channel: str, src: str, expires=3600):
        # expires is in seconds, so the default is 60 minutes!
        with self._get_fs() as fs:
            expiration_timestamp = (
                int(datetime.datetime.now(tz=datetime.timezone.utc).timestamp())
                + expires
            )
            redirect_url = fs.sign(
                path.join(self._bucket_map(channel), src), expiration_timestamp
            )
            return redirect_url

    def get_filemetadata(self, channel: str, src: str):
        with self._get_fs() as fs:
            filepath = path.join(self._bucket_map(channel), src)
            infodata = fs.info(filepath)

            if infodata['type'] != 'file':
                raise FileNotFoundError()

            mtime = datetime.datetime.fromisoformat(
                infodata['updated'].replace('Z', '+00:00')
            ).timestamp()
            msize = infodata['size']
            etag = infodata['etag']

            return (msize, mtime, etag)

    def cleanup_temp_files(self, channel: str, dry_run: bool = False):
        with self._get_fs() as fs:
            temp_files = []
            for each_end in [".bz2", ".gz"]:
                temp_files.extend(
                    fs.glob(f"{self._bucket_map(channel)}/**/*.json?*{each_end}")
                )

            for each_temp_file in temp_files:
                logger.info(f"removing {each_temp_file} from pkgstore")
                if not dry_run:
                    fs.delete(each_temp_file)
