#
# Copyright (c) 2012-2021 Snowflake Computing Inc. All right reserved.
#

from __future__ import division

import json
import os
from collections import namedtuple
from logging import getLogger
from typing import TYPE_CHECKING, Any, Dict

from azure.core.exceptions import HttpResponseError, ResourceNotFoundError
from azure.storage.blob import BlobServiceClient, ContentSettings, ExponentialRetry

from .constants import HTTP_HEADER_VALUE_OCTET_STREAM, FileHeader, ResultStatus
from .encryption_util import EncryptionMetadata

if TYPE_CHECKING:  # pragma: no cover
    from .file_transfer_agent_sdk import SnowflakeFileMeta

logger = getLogger(__name__)

"""
Azure Location: Azure container name + path
"""
AzureLocation = namedtuple(
    "AzureLocation",
    ["container_name", "path"],  # Azure container name  # Azure path name
)


class SnowflakeAzureUtil(object):
    """Azure Utility class."""

    @staticmethod
    def create_client(
        stage_info: Dict[str, Any],
        use_accelerate_endpoint: bool = False,
        use_s3_regional_url: bool = False,
        s3_connection_pool_size: int = 1,
    ) -> BlobServiceClient:
        """Creates a client object with a stage credential.

        Args:
            stage_info: Information about the stage.
            use_accelerate_endpoint: Not used for Azure client.

        Returns:
            The client to communicate with GCS.
        """
        stage_credentials = stage_info["creds"]
        sas_token = stage_credentials["AZURE_SAS_TOKEN"]
        if sas_token and sas_token.startswith("?"):
            sas_token = sas_token[1:]
        end_point = stage_info["endPoint"]
        if end_point.startswith("blob."):
            end_point = end_point[len("blob.") :]
        client = BlobServiceClient(
            account_url=f"https://{stage_info['storageAccount']}.blob.{end_point}",
            credential=sas_token,
        )
        client._config.retry_policy = ExponentialRetry(
            initial_backoff=1, increment_base=2, max_attempts=60, random_jitter_range=2
        )

        return client

    @staticmethod
    def extract_container_name_and_path(stage_location):
        stage_location = os.path.expanduser(stage_location)
        container_name = stage_location
        path = ""

        # split stage location as bucket name and path
        if "/" in stage_location:
            container_name = stage_location[0 : stage_location.index("/")]
            path = stage_location[stage_location.index("/") + 1 :]
            if path and not path.endswith("/"):
                path += "/"

        return AzureLocation(container_name=container_name, path=path)

    @staticmethod
    def get_file_header(meta: "SnowflakeFileMeta", filename):
        """Gets Azure file properties."""
        client: BlobServiceClient = meta.client_meta.cloud_client
        azure_location = SnowflakeAzureUtil.extract_container_name_and_path(
            meta.client_meta.stage_info["location"]
        )
        try:
            # HTTP HEAD request
            blob = client.get_blob_client(
                azure_location.container_name, azure_location.path + filename
            )
            blob_details = blob.get_blob_properties()
        except ResourceNotFoundError:
            meta.result_status = ResultStatus.NOT_FOUND_FILE
            return FileHeader(
                digest=None, content_length=None, encryption_metadata=None
            )
        except HttpResponseError as err:
            logger.debug(
                f"Caught exception's status code: {err.status_code} and message: {str(err)}"
            )
            if (
                err.status_code == 403
                and SnowflakeAzureUtil._detect_azure_token_expire_error(err)
            ):
                logger.debug("AZURE Token expired. Renew and retry")
                meta.result_status = ResultStatus.RENEW_TOKEN
            else:
                logger.debug(
                    f"Unexpected Azure error: {err} "
                    f"container: {azure_location.container_name}, path: {azure_location.path}"
                )
                meta.result_status = ResultStatus.ERROR

            return
        meta.result_status = ResultStatus.UPLOADED
        encryptiondata = json.loads(blob_details.metadata.get("encryptiondata", "null"))
        encryption_metadata = (
            EncryptionMetadata(
                key=encryptiondata["WrappedContentKey"]["EncryptedKey"],
                iv=encryptiondata["ContentEncryptionIV"],
                matdesc=blob_details.metadata["matdesc"],
            )
            if encryptiondata
            else None
        )

        return FileHeader(
            digest=blob_details.metadata.get("sfcdigest"),
            content_length=blob_details.size,
            encryption_metadata=encryption_metadata,
        )

    @staticmethod
    def _detect_azure_token_expire_error(err):
        if err.status_code != 403:
            return False
        errstr = str(err)
        return (
            "Signature not valid in the specified time frame" in errstr
            or "Server failed to authenticate the request." in errstr
        )

    @staticmethod
    def upload_file(
        data_file: str,
        meta: "SnowflakeFileMeta",
        encryption_metadata: "EncryptionMetadata",
        max_concurrency: int,
        multipart_threshold: int,
    ):
        """Uploads the local file to Azure's Blob Storage.

        Args:
            data_file: File path on local system.
            meta: The File meta object (contains credentials and remote location).
            encryption_metadata: Encryption metadata to be set on object.
            max_concurrency: Not applicable to Azure.
            multipart_threshold: The number of bytes after which size a file should be uploaded concurrently in chunks.
                Not applicable to Azure.

        Raises:
            HTTPError if some http errors occurred.

        Returns:
            None.
        """
        azure_metadata = {
            "sfcdigest": meta.sha256_digest,
        }
        if encryption_metadata:
            azure_metadata.update(
                {
                    "encryptiondata": json.dumps(
                        {
                            "EncryptionMode": "FullBlob",
                            "WrappedContentKey": {
                                "KeyId": "symmKey1",
                                "EncryptedKey": encryption_metadata.key,
                                "Algorithm": "AES_CBC_256",
                            },
                            "EncryptionAgent": {
                                "Protocol": "1.0",
                                "EncryptionAlgorithm": "AES_CBC_128",
                            },
                            "ContentEncryptionIV": encryption_metadata.iv,
                            "KeyWrappingMetadata": {"EncryptionLibrary": "Java 5.3.0"},
                        }
                    ),
                    "matdesc": encryption_metadata.matdesc,
                }
            )
        azure_location = SnowflakeAzureUtil.extract_container_name_and_path(
            meta.client_meta.stage_info["location"]
        )
        path = azure_location.path + meta.dst_file_name.lstrip("/")

        client: BlobServiceClient = meta.client_meta.cloud_client
        callback = None
        upload_src = None
        upload_size = None

        if meta.src_stream is None:
            upload_size = os.path.getsize(data_file)
            upload_src = open(data_file, "rb")
        else:
            upload_src = meta.real_src_stream or meta.src_stream
            upload_size = upload_src.seek(0, os.SEEK_END)
            upload_src.seek(0)

        if meta.put_azure_callback:
            callback = meta.put_azure_callback(
                data_file,
                upload_size,
                output_stream=meta.put_callback_output_stream,
                show_progress_bar=meta.show_progress_bar,
            )

        def azure_callback(response):
            current = response.context["upload_stream_current"]
            total = response.context["data_stream_total"]
            if current is not None:
                callback(current)
                logger.debug(
                    "data transfer progress from sdk callback. "
                    f"current: {current}, total: {total}"
                )

        try:
            blob = client.get_blob_client(azure_location.container_name, path)
            blob.upload_blob(
                upload_src,
                metadata=azure_metadata,
                overwrite=True,
                max_concurrency=max_concurrency,
                raw_response_hook=azure_callback if meta.put_azure_callback else None,
                content_settings=ContentSettings(
                    content_type=HTTP_HEADER_VALUE_OCTET_STREAM,
                    content_encoding="utf-8",
                ),
            )
        except HttpResponseError as err:
            logger.debug(
                f"Caught exception's status code: {err.status_code} and message: {err}"
            )
            if (
                err.status_code == 403
                and SnowflakeAzureUtil._detect_azure_token_expire_error(err)
            ):
                logger.debug("AZURE Token expired. Renew and retry")
                meta.result_status = ResultStatus.RENEW_TOKEN
            else:
                meta.last_error = err
                meta.result_status = ResultStatus.NEED_RETRY
            return
        finally:
            if meta.src_stream is None:
                upload_src.close()

        logger.debug("DONE putting a file")
        meta.dst_file_size = meta.upload_size
        meta.result_status = ResultStatus.UPLOADED
        # Comparing with s3, azure haven't experienced OpenSSL.SSL.SysCallError,
        # so we will add logic to catch it only when it happens

    @staticmethod
    def _native_download_file(
        meta: "SnowflakeFileMeta", full_dst_file_name, max_concurrency
    ):
        azure_location = SnowflakeAzureUtil.extract_container_name_and_path(
            meta.client_meta.stage_info["location"]
        )
        path = azure_location.path + meta.src_file_name.lstrip("/")
        client: BlobServiceClient = meta.client_meta.cloud_client

        callback = None
        if meta.get_azure_callback:
            callback = meta.get_azure_callback(
                meta.src_file_name,
                meta.src_file_size,
                output_stream=meta.get_callback_output_stream,
                show_progress_bar=meta.show_progress_bar,
            )

        def azure_callback(response):
            current = response.context["download_stream_current"]
            total = response.context["data_stream_total"]
            if current is not None:
                callback(current)
                logger.debug(
                    f"data transfer progress from sdk callback. current: {current}, total: {total}"
                )

        try:
            blob = client.get_blob_client(azure_location.container_name, path)
            with open(full_dst_file_name, "wb") as download_f:
                download = blob.download_blob(
                    max_concurrency=max_concurrency,
                    raw_response_hook=azure_callback
                    if meta.put_azure_callback
                    else None,
                )
                download.readinto(download_f)

        except HttpResponseError as err:
            logger.debug(
                f"Caught exception's status code: {err.status_code} and message: {str(err)}"
            )
            if (
                err.status_code == 403
                and SnowflakeAzureUtil._detect_azure_token_expire_error(err)
            ):
                logger.debug("AZURE Token expired. Renew and retry")
                meta.result_status = ResultStatus.RENEW_TOKEN
            else:
                meta.last_error = err
                meta.result_status = ResultStatus.NEED_RETRY
            return
        meta.result_status = ResultStatus.DOWNLOADED
