# -*- 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

[((54, 27, 56, 91), 'polystores.clients.azure_client.get_blob_service_connection', 'get_blob_service_connection', (), '', False, 'from polystores.clients.azure_client import get_blob_service_connection\n'), ((187, 18, 187, 42), 'os.path.abspath', 'os.path.abspath', ({(187, 34, 187, 41): 'dirname'}, {}), '(dirname)', False, 'import os\n'), ((209, 21, 209, 48), 'os.path.abspath', 'os.path.abspath', ({(209, 37, 209, 47): 'local_path'}, {}), '(local_path)', False, 'import os\n'), ((214, 8, 214, 40), 'polystores.utils.check_dirname_exists', 'check_dirname_exists', ({(214, 29, 214, 39): 'local_path'}, {}), '(local_path)', False, 'from polystores.utils import append_basename, check_dirname_exists, get_files_in_current_directory\n'), ((234, 21, 234, 48), 'os.path.abspath', 'os.path.abspath', ({(234, 37, 234, 47): 'local_path'}, {}), '(local_path)', False, 'import os\n'), ((75, 19, 75, 58), 'rhea.parser.parse_wasbs_path', 'rhea_parser.parse_wasbs_path', ({(75, 48, 75, 57): 'wasbs_url'}, {}), '(wasbs_url)', True, 'from rhea import parser as rhea_parser\n'), ((124, 21, 124, 47), 'os.path.join', 'os.path.join', ({(124, 34, 124, 40): 'prefix', (124, 42, 124, 46): 'path'}, {}), '(prefix, path)', False, 'import os\n'), ((166, 19, 166, 50), 'polystores.utils.append_basename', 'append_basename', ({(166, 35, 166, 39): 'blob', (166, 41, 166, 49): 'filename'}, {}), '(blob, filename)', False, 'from polystores.utils import append_basename, check_dirname_exists, get_files_in_current_directory\n'), ((184, 19, 184, 49), 'polystores.utils.append_basename', 'append_basename', ({(184, 35, 184, 39): 'blob', (184, 41, 184, 48): 'dirname'}, {}), '(blob, dirname)', False, 'from polystores.utils import append_basename, check_dirname_exists, get_files_in_current_directory\n'), ((188, 13, 188, 52), 'polystores.utils.get_files_in_current_directory', 'get_files_in_current_directory', ({(188, 44, 188, 51): 'dirname'}, {}), '(dirname)', False, 'from polystores.utils import append_basename, check_dirname_exists, get_files_in_current_directory\n'), ((212, 25, 212, 58), 'polystores.utils.append_basename', 'append_basename', ({(212, 41, 212, 51): 'local_path', (212, 53, 212, 57): 'blob'}, {}), '(local_path, blob)', False, 'from polystores.utils import append_basename, check_dirname_exists, get_files_in_current_directory\n'), ((237, 25, 237, 58), 'polystores.utils.append_basename', 'append_basename', ({(237, 41, 237, 51): 'local_path', (237, 53, 237, 57): 'blob'}, {}), '(local_path, blob)', False, 'from polystores.utils import append_basename, check_dirname_exists, get_files_in_current_directory\n'), ((240, 12, 240, 57), 'polystores.utils.check_dirname_exists', 'check_dirname_exists', (), '', False, 'from polystores.utils import append_basename, check_dirname_exists, get_files_in_current_directory\n'), ((248, 23, 248, 55), 'os.path.join', 'os.path.join', ({(248, 36, 248, 46): 'local_path', (248, 48, 248, 54): 'prefix'}, {}), '(local_path, prefix)', False, 'import os\n'), ((249, 21, 249, 47), 'os.path.join', 'os.path.join', ({(249, 34, 249, 38): 'blob', (249, 40, 249, 46): 'prefix'}, {}), '(blob, prefix)', False, 'import os\n'), ((259, 23, 259, 57), 'os.path.join', 'os.path.join', ({(259, 36, 259, 46): 'local_path', (259, 48, 259, 56): 'file_key'}, {}), '(local_path, file_key)', False, 'import os\n'), ((260, 23, 260, 51), 'os.path.join', 'os.path.join', ({(260, 36, 260, 40): 'blob', (260, 42, 260, 50): 'file_key'}, {}), '(blob, file_key)', False, 'import os\n'), ((277, 21, 277, 47), 'os.path.join', 'os.path.join', ({(277, 34, 277, 38): 'blob', (277, 40, 277, 46): 'prefix'}, {}), '(blob, prefix)', False, 'import os\n'), ((284, 23, 284, 51), 'os.path.join', 'os.path.join', ({(284, 36, 284, 40): 'blob', (284, 42, 284, 50): 'file_key'}, {}), '(blob, file_key)', False, 'import os\n'), ((78, 18, 78, 44), 'polystores.exceptions.PolyaxonStoresException', 'PolyaxonStoresException', ({(78, 42, 78, 43): 'e'}, {}), '(e)', False, 'from polystores.exceptions import PolyaxonStoresException\n'), ((219, 18, 219, 44), 'polystores.exceptions.PolyaxonStoresException', 'PolyaxonStoresException', ({(219, 42, 219, 43): 'e'}, {}), '(e)', False, 'from polystores.exceptions import PolyaxonStoresException\n'), ((242, 12, 242, 35), 'os.makedirs', 'os.makedirs', ({(242, 24, 242, 34): 'local_path'}, {}), '(local_path)', False, 'import os\n'), ((190, 47, 190, 74), 'os.path.relpath', 'os.path.relpath', ({(190, 63, 190, 64): 'f', (190, 66, 190, 73): 'dirname'}, {}), '(f, dirname)', False, 'import os\n')]