# -*- coding: utf-8 -*-
from __future__ import absolute_import, division, print_function

import os

from rhea import RheaError
from rhea import parser as rhea_parser

from azure.common import AzureHttpError
from azure.storage.blob.models import BlobPrefix

from polystores.clients.azure_client import get_blob_service_connection
from polystores.exceptions import PolyaxonStoresException
from polystores.stores.base_store import BaseStore
from polystores.utils import append_basename, check_dirname_exists, get_files_in_current_directory

# pylint:disable=arguments-differ


class AzureStore(BaseStore):
    """
    Azure store Service.
    """
    STORE_TYPE = BaseStore._AZURE_STORE  # pylint:disable=protected-access

    def __init__(self, connection=None, **kwargs):
        self._connection = connection
        self._account_name = kwargs.get('account_name') or kwargs.get('AZURE_ACCOUNT_NAME')
        self._account_key = kwargs.get('account_key') or kwargs.get('AZURE_ACCOUNT_KEY')
        self._connection_string = (
            kwargs.get('connection_string') or kwargs.get('AZURE_CONNECTION_STRING'))

    @property
    def connection(self):
        if self._connection is None:
            self.set_connection(account_name=self._account_name,
                                account_key=self._account_key,
                                connection_string=self._connection_string)
        return self._connection

    def set_connection(self, account_name=None, account_key=None, connection_string=None):
        """
        Sets a new Blob service connection.

        Args:
        account_name: `str`. The storage account name.
        account_key: `str`. The storage account key.
        connection_string: `str`. If specified, this will override all other parameters besides
        request session.

        Returns:
            BlockBlobService instance
        """
        self._connection = get_blob_service_connection(account_name=account_name,
                                                       account_key=account_key,
                                                       connection_string=connection_string)

    def set_env_vars(self):
        if self._account_name:
            os.environ['AZURE_ACCOUNT_NAME'] = self._account_name
        if self._account_key:
            os.environ['AZURE_ACCOUNT_KEY'] = self._account_key
        if self._connection_string:
            os.environ['AZURE_CONNECTION_STRING'] = self._connection_string

    @staticmethod
    def parse_wasbs_url(wasbs_url):
        """
        Parses and validates a wasbs url.

        Returns:
            tuple(container, storage_account, path).
        """
        try:
            spec = rhea_parser.parse_wasbs_path(wasbs_url)
            return spec.container, spec.storage_account, spec.path
        except RheaError as e:
            raise PolyaxonStoresException(e)

    def check_blob(self, blob, container_name=None):
        """
        Checks if a blob exists.

        Args:
            blob: `str`. Name of existing blob.
            container_name: `str`. Name of existing container.

        Returns:
            bool
        """
        if not container_name:
            container_name, _, blob = self.parse_wasbs_url(blob)
        try:
            return self.connection.get_blob_properties(
                container_name,
                blob
            )
        except AzureHttpError:
            return None

    def ls(self, path):
        results = self.list(key=path)
        return {'files': results['blobs'], 'dirs': results['prefixes']}

    def list(self, key, container_name=None, path=None, delimiter='/', marker=None):
        """
        Checks if a blob exists.

        Args:
            key: `str`. key prefix.
            container_name: `str`. Name of existing container.
            path: `str`. an extra path to append to the key.
            delimiter: `str`. the delimiter marks key hierarchy.
            marker: `str`. An opaque continuation token.
        """
        if not container_name:
            container_name, _, key = self.parse_wasbs_url(key)

        if key and not key.endswith('/'):
            key += '/'

        prefix = key
        if path:
            prefix = os.path.join(prefix, path)

        if prefix and not prefix.endswith('/'):
            prefix += '/'

        list_blobs = []
        list_prefixes = []
        while True:
            results = self.connection.list_blobs(container_name,
                                                 prefix=prefix,
                                                 delimiter=delimiter,
                                                 marker=marker)
            for r in results:
                if isinstance(r, BlobPrefix):
                    name = r.name[len(key):]
                    list_prefixes.append(name)
                else:
                    name = r.name[len(key):]
                    list_blobs.append((name, r.properties.content_length))
            if results.next_marker:
                marker = results.next_marker
            else:
                break
        return {
            'blobs': list_blobs,
            'prefixes': list_prefixes
        }

    def upload_file(self, filename, blob, container_name=None, use_basename=True):
        """
        Uploads a local file to Google Cloud Storage.

        Args:
            filename: `str`. the file to upload.
            blob: `str`. blob to upload to.
            container_name: `str`. the name of the container.
            use_basename: `bool`. whether or not to use the basename of the filename.
        """
        if not container_name:
            container_name, _, blob = self.parse_wasbs_url(blob)

        if use_basename:
            blob = append_basename(blob, filename)

        self.connection.create_blob_from_path(container_name, blob, filename)

    def upload_dir(self, dirname, blob, container_name=None, use_basename=True):
        """
        Uploads a local directory to to Google Cloud Storage.

        Args:
            dirname: `str`. name of the directory to upload.
            blob: `str`. blob to upload to.
            container_name: `str`. the name of the container.
            use_basename: `bool`. whether or not to use the basename of the directory.
        """
        if not container_name:
            container_name, _, blob = self.parse_wasbs_url(blob)

        if use_basename:
            blob = append_basename(blob, dirname)

        # Turn the path to absolute paths
        dirname = os.path.abspath(dirname)
        with get_files_in_current_directory(dirname) as files:
            for f in files:
                file_blob = os.path.join(blob, os.path.relpath(f, dirname))
                self.upload_file(filename=f,
                                 blob=file_blob,
                                 container_name=container_name,
                                 use_basename=False)

    def download_file(self, blob, local_path, container_name=None, use_basename=True):
        """
        Downloads a file from Google Cloud Storage.

        Args:
            blob: `str`. blob to download.
            local_path: `str`. the path to download to.
            container_name: `str`. the name of the container.
            use_basename: `bool`. whether or not to use the basename of the blob.
        """
        if not container_name:
            container_name, _, blob = self.parse_wasbs_url(blob)

        local_path = os.path.abspath(local_path)

        if use_basename:
            local_path = append_basename(local_path, blob)

        check_dirname_exists(local_path)

        try:
            self.connection.get_blob_to_path(container_name, blob, local_path)
        except AzureHttpError as e:
            raise PolyaxonStoresException(e)

    def download_dir(self, blob, local_path, container_name=None, use_basename=True):
        """
        Download a directory from Google Cloud Storage.

        Args:
            blob: `str`. blob to download.
            local_path: `str`. the path to download to.
            container_name: `str`. the name of the container.
            use_basename: `bool`. whether or not to use the basename of the key.
        """
        if not container_name:
            container_name, _, blob = self.parse_wasbs_url(blob)

        local_path = os.path.abspath(local_path)

        if use_basename:
            local_path = append_basename(local_path, blob)

        try:
            check_dirname_exists(local_path, is_dir=True)
        except PolyaxonStoresException:
            os.makedirs(local_path)

        results = self.list(container_name=container_name, key=blob, delimiter='/')

        # Create directories
        for prefix in sorted(results['prefixes']):
            direname = os.path.join(local_path, prefix)
            prefix = os.path.join(blob, prefix)
            # Download files under
            self.download_dir(blob=prefix,
                              local_path=direname,
                              container_name=container_name,
                              use_basename=False)

        # Download files
        for file_key in results['blobs']:
            file_key = file_key[0]
            filename = os.path.join(local_path, file_key)
            file_key = os.path.join(blob, file_key)
            self.download_file(blob=file_key,
                               local_path=filename,
                               container_name=container_name,
                               use_basename=False)

    def delete(self, blob, container_name=None):
        if not container_name:
            container_name, _, blob = self.parse_wasbs_url(blob)

        results = self.list(container_name=container_name, key=blob, delimiter='/')

        if not any([results['prefixes'], results['blobs']]):
            self.delete_file(blob=blob, container_name=container_name)

        # Delete directories
        for prefix in sorted(results['prefixes']):
            prefix = os.path.join(blob, prefix)
            # Download files under
            self.delete(blob=prefix, container_name=container_name)

        # Delete files
        for file_key in results['blobs']:
            file_key = file_key[0]
            file_key = os.path.join(blob, file_key)
            self.delete_file(blob=file_key, container_name=container_name)

    def delete_file(self, blob, container_name=None):
        """
        Deletes if a blob exists.

        Args:
            blob: `str`. Name of existing blob.
            container_name: `str`. Name of existing container.
        """
        if not container_name:
            container_name, _, blob = self.parse_wasbs_url(blob)

        try:
            self.connection.delete_blob(container_name, blob)
        except AzureHttpError:
            pass

[((1913, 2034), 'polystores.clients.azure_client.get_blob_service_connection', 'get_blob_service_connection', ([], {'account_name': 'account_name', 'account_key': 'account_key', 'connection_string': 'connection_string'}), '(account_name=account_name, account_key=\n    account_key, connection_string=connection_string)\n', (1940, 2034), False, 'from polystores.clients.azure_client import get_blob_service_connection\n'), ((6539, 6563), 'os.path.abspath', 'os.path.abspath', (['dirname'], {}), '(dirname)\n', (6554, 6563), False, 'import os\n'), ((7482, 7509), 'os.path.abspath', 'os.path.abspath', (['local_path'], {}), '(local_path)\n', (7497, 7509), False, 'import os\n'), ((7604, 7636), 'polystores.utils.check_dirname_exists', 'check_dirname_exists', (['local_path'], {}), '(local_path)\n', (7624, 7636), False, 'from polystores.utils import append_basename, check_dirname_exists, get_files_in_current_directory\n'), ((8353, 8380), 'os.path.abspath', 'os.path.abspath', (['local_path'], {}), '(local_path)\n', (8368, 8380), False, 'import os\n'), ((2696, 2735), 'rhea.parser.parse_wasbs_path', 'rhea_parser.parse_wasbs_path', (['wasbs_url'], {}), '(wasbs_url)\n', (2724, 2735), True, 'from rhea import parser as rhea_parser\n'), ((4233, 4259), 'os.path.join', 'os.path.join', (['prefix', 'path'], {}), '(prefix, path)\n', (4245, 4259), False, 'import os\n'), ((5758, 5789), 'polystores.utils.append_basename', 'append_basename', (['blob', 'filename'], {}), '(blob, filename)\n', (5773, 5789), False, 'from polystores.utils import append_basename, check_dirname_exists, get_files_in_current_directory\n'), ((6447, 6477), 'polystores.utils.append_basename', 'append_basename', (['blob', 'dirname'], {}), '(blob, dirname)\n', (6462, 6477), False, 'from polystores.utils import append_basename, check_dirname_exists, get_files_in_current_directory\n'), ((6577, 6616), 'polystores.utils.get_files_in_current_directory', 'get_files_in_current_directory', (['dirname'], {}), '(dirname)\n', (6607, 6616), False, 'from polystores.utils import append_basename, check_dirname_exists, get_files_in_current_directory\n'), ((7561, 7594), 'polystores.utils.append_basename', 'append_basename', (['local_path', 'blob'], {}), '(local_path, blob)\n', (7576, 7594), False, 'from polystores.utils import append_basename, check_dirname_exists, get_files_in_current_directory\n'), ((8432, 8465), 'polystores.utils.append_basename', 'append_basename', (['local_path', 'blob'], {}), '(local_path, blob)\n', (8447, 8465), False, 'from polystores.utils import append_basename, check_dirname_exists, get_files_in_current_directory\n'), ((8492, 8537), 'polystores.utils.check_dirname_exists', 'check_dirname_exists', (['local_path'], {'is_dir': '(True)'}), '(local_path, is_dir=True)\n', (8512, 8537), False, 'from polystores.utils import append_basename, check_dirname_exists, get_files_in_current_directory\n'), ((8803, 8835), 'os.path.join', 'os.path.join', (['local_path', 'prefix'], {}), '(local_path, prefix)\n', (8815, 8835), False, 'import os\n'), ((8857, 8883), 'os.path.join', 'os.path.join', (['blob', 'prefix'], {}), '(blob, prefix)\n', (8869, 8883), False, 'import os\n'), ((9250, 9284), 'os.path.join', 'os.path.join', (['local_path', 'file_key'], {}), '(local_path, file_key)\n', (9262, 9284), False, 'import os\n'), ((9308, 9336), 'os.path.join', 'os.path.join', (['blob', 'file_key'], {}), '(blob, file_key)\n', (9320, 9336), False, 'import os\n'), ((10014, 10040), 'os.path.join', 'os.path.join', (['blob', 'prefix'], {}), '(blob, prefix)\n', (10026, 10040), False, 'import os\n'), ((10268, 10296), 'os.path.join', 'os.path.join', (['blob', 'file_key'], {}), '(blob, file_key)\n', (10280, 10296), False, 'import os\n'), ((2852, 2878), 'polystores.exceptions.PolyaxonStoresException', 'PolyaxonStoresException', (['e'], {}), '(e)\n', (2875, 2878), False, 'from polystores.exceptions import PolyaxonStoresException\n'), ((7784, 7810), 'polystores.exceptions.PolyaxonStoresException', 'PolyaxonStoresException', (['e'], {}), '(e)\n', (7807, 7810), False, 'from polystores.exceptions import PolyaxonStoresException\n'), ((8590, 8613), 'os.makedirs', 'os.makedirs', (['local_path'], {}), '(local_path)\n', (8601, 8613), False, 'import os\n'), ((6702, 6729), 'os.path.relpath', 'os.path.relpath', (['f', 'dirname'], {}), '(f, dirname)\n', (6717, 6729), False, 'import os\n')]