# ------------------------------------
# Copyright (c) Microsoft Corporation.
# Licensed under the MIT License.
# ------------------------------------
import base64
import functools
import json
from typing import Any, Callable, Optional, overload
from urllib.parse import urlparse

from typing_extensions import Literal

from azure.core.pipeline import PipelineResponse
from azure.core.polling import LROPoller
from azure.core.tracing.decorator import distributed_trace

from ._generated.models import PreBackupOperationParameters, PreRestoreOperationParameters, SASTokenParameter
from ._models import KeyVaultBackupResult
from ._internal import KeyVaultClientBase, parse_folder_url
from ._internal.polling import KeyVaultBackupClientPolling, KeyVaultBackupClientPollingMethod


def _parse_status_url(url: str) -> str:
    parsed = urlparse(url)
    job_id = parsed.path.split("/")[2]
    return job_id


def _get_continuation_token(pipeline_response: PipelineResponse) -> str:
    """Returns an opaque token which can be used by the user to rehydrate/restart the LRO.

    Saves the state of the LRO based on a status response so that polling can be resumed from that context. Because
    the service has different operations for backup/restore starting vs. status checking, the caller is expected to
    first use the status URL from the initial response to make a status request and then pass that status response to
    this function to be serialized into a continuation token.

    :param pipeline_response: The pipeline response of the operation status request.
    :type pipeline_response: ~azure.core.pipeline.PipelineResponse
    :returns: An opaque continuation token that can be provided to Core to rehydrate the LRO.
    :rtype: str
    """
    # Headers needed for LRO rehydration - use an allowlist approach for security
    lro_headers = {"azure-asyncoperation", "operation-location", "location", "content-type", "retry-after"}
    response = pipeline_response.http_response
    filtered_headers = {k: v for k, v in response.headers.items() if k.lower() in lro_headers}

    request = response.request
    # Serialize the essential parts of the PipelineResponse to JSON.
    if request:
        request_headers = {}
        # Preserve x-ms-client-request-id for request correlation
        if "x-ms-client-request-id" in request.headers:
            request_headers["x-ms-client-request-id"] = request.headers["x-ms-client-request-id"]
        request_state = {
            "method": request.method,
            "url": request.url,
            "headers": request_headers,
        }
    else:
        request_state = None

    # Use a versioned token schema: {"version": <int>, "data": <your state>}
    # This allows for future compatibility checking when deserializing
    token = {
        "version": 1,
        "data": {
            "request": request_state,
            "response": {
                "status_code": response.status_code,
                "headers": filtered_headers,
                "content": base64.b64encode(response.content).decode("ascii"),
            },
        },
    }
    return base64.b64encode(json.dumps(token, separators=(",", ":")).encode("utf-8")).decode("ascii")


class KeyVaultBackupClient(KeyVaultClientBase):
    """Performs Key Vault backup and restore operations.

    :param str vault_url: URL of the vault on which the client will operate. This is also called the vault's "DNS Name".
        You should validate that this URL references a valid Key Vault or Managed HSM resource.
        See https://aka.ms/azsdk/blog/vault-uri for details.
    :param credential: An object which can provide an access token for the vault, such as a credential from
        :mod:`azure.identity`
    :type credential: ~azure.core.credentials.TokenCredential

    :keyword api_version: Version of the service API to use. Defaults to the most recent.
    :paramtype api_version: ~azure.keyvault.administration.ApiVersion or str
    :keyword bool verify_challenge_resource: Whether to verify the authentication challenge resource matches the Key
        Vault or Managed HSM domain. Defaults to True.
    """

    def _use_continuation_token(self, continuation_token: str, status_method: Callable) -> str:
        status_url = base64.b64decode(continuation_token.encode()).decode("ascii")
        try:
            job_id = _parse_status_url(status_url)
        except Exception as ex:  # pylint: disable=broad-except
            raise ValueError(
                "The provided continuation_token is malformed. A valid token can be obtained from the "
                + "operation poller's continuation_token() method"
            ) from ex

        pipeline_response = status_method(job_id=job_id, cls=lambda pipeline_response, _, __: pipeline_response)
        if "azure-asyncoperation" not in pipeline_response.http_response.headers:
            pipeline_response.http_response.headers["azure-asyncoperation"] = status_url
        return _get_continuation_token(pipeline_response)

    @overload
    def begin_backup(
        self,
        blob_storage_url: str,
        *,
        use_managed_identity: Literal[True],
        continuation_token: Optional[str] = None,
        **kwargs: Any,
    ) -> LROPoller[KeyVaultBackupResult]: ...

    @overload
    def begin_backup(
        self,
        blob_storage_url: str,
        *,
        sas_token: str,
        continuation_token: Optional[str] = None,
        **kwargs: Any,
    ) -> LROPoller[KeyVaultBackupResult]: ...

    # Disabling pylint checks because they don't correctly handle overloads
    @distributed_trace
    def begin_backup(  # pylint: disable=docstring-missing-param,docstring-keyword-should-match-keyword-only
        self, blob_storage_url: str, *args: str, **kwargs: Any
    ) -> LROPoller[KeyVaultBackupResult]:
        """Begin a full backup of the Key Vault.

        :param str blob_storage_url: URL of the blob storage container in which the backup will be stored, for example
            https://<account>.blob.core.windows.net/backup.

        :keyword str sas_token: Optional Shared Access Signature (SAS) token to authorize access to the blob. Required
            unless `use_managed_identity` is set to True.
        :keyword use_managed_identity: Indicates which authentication method should be used. If set to True, Managed HSM
            will use the configured user-assigned managed identity to authenticate with Azure Storage. Otherwise, a SAS
            token has to be specified.
        :paramtype use_managed_identity: bool
        :keyword str continuation_token: A continuation token to restart polling from a saved state.

        :returns: An :class:`~azure.core.polling.LROPoller` instance. Call `result()` on this object to wait for the
            operation to complete and get a :class:`KeyVaultBackupResult`.
        :rtype: ~azure.core.polling.LROPoller[~azure.keyvault.administration.KeyVaultBackupResult]

        Example:
            .. literalinclude:: ../tests/test_examples_administration.py
                :start-after: [START begin_backup]
                :end-before: [END begin_backup]
                :language: python
                :caption: Create a vault backup
                :dedent: 8
        """
        polling_interval = kwargs.pop("_polling_interval", 5)
        continuation_token = kwargs.pop("continuation_token", None)
        use_managed_identity = kwargs.pop("use_managed_identity", False)
        # `sas_token` was formerly a required positional parameter
        try:
            sas_token: Optional[str] = args[0]
        except IndexError:
            sas_token = kwargs.pop("sas_token", None)
        sas_parameter = self._models.SASTokenParameter(
            storage_resource_uri=blob_storage_url, token=sas_token, use_managed_identity=use_managed_identity
        )

        status_response = None
        if continuation_token:
            status_response = self._use_continuation_token(continuation_token, self._client.full_backup_status)

        return self._client.begin_full_backup(
            azure_storage_blob_container_uri=sas_parameter,
            cls=KeyVaultBackupResult._from_generated,  # pylint: disable=protected-access
            continuation_token=status_response,
            polling=KeyVaultBackupClientPollingMethod(
                lro_algorithms=[KeyVaultBackupClientPolling()], timeout=polling_interval, **kwargs
            ),
            **kwargs,
        )

    @overload
    def begin_restore(
        self,
        folder_url: str,
        *,
        use_managed_identity: Literal[True],
        key_name: Optional[str] = None,
        continuation_token: Optional[str] = None,
        **kwargs: Any,
    ) -> LROPoller[None]: ...

    @overload
    def begin_restore(
        self,
        folder_url: str,
        *,
        sas_token: str,
        key_name: Optional[str] = None,
        continuation_token: Optional[str] = None,
        **kwargs: Any,
    ) -> LROPoller[None]: ...

    # Disabling pylint checks because they don't correctly handle overloads
    @distributed_trace
    def begin_restore(  # pylint: disable=docstring-missing-param,docstring-keyword-should-match-keyword-only
        self, folder_url: str, *args: str, **kwargs: Any
    ) -> LROPoller[None]:
        """Restore a Key Vault backup.

        This method restores either a complete Key Vault backup or when ``key_name`` has a value, a single key.

        :param str folder_url: URL of the blob holding the backup. This would be the `folder_url` of a
            :class:`KeyVaultBackupResult` returned by :func:`begin_backup`, for example
            https://<account>.blob.core.windows.net/backup/mhsm-account-2020090117323313

        :keyword str sas_token: Optional Shared Access Signature (SAS) token to authorize access to the blob. Required
            unless `use_managed_identity` is set to True.
        :keyword use_managed_identity: Indicates which authentication method should be used. If set to True, Managed HSM
            will use the configured user-assigned managed identity to authenticate with Azure Storage. Otherwise, a SAS
            token has to be specified.
        :paramtype use_managed_identity: bool
        :keyword str key_name: Name of a single key in the backup. When set, only this key will be restored.
        :keyword str continuation_token: A continuation token to restart polling from a saved state.

        :returns: An :class:`~azure.core.polling.LROPoller` instance. Call `wait()` or `result()` on this object to wait
            for the operation to complete (the return value is None in either case).
        :rtype: ~azure.core.polling.LROPoller

        Examples:
            .. literalinclude:: ../tests/test_examples_administration.py
                :start-after: [START begin_restore]
                :end-before: [END begin_restore]
                :language: python
                :caption: Restore a vault backup
                :dedent: 8

            .. literalinclude:: ../tests/test_examples_administration.py
                :start-after: [START begin_selective_restore]
                :end-before: [END begin_selective_restore]
                :language: python
                :caption: Restore a single key
                :dedent: 8
        """
        polling_interval = kwargs.pop("_polling_interval", 5)
        continuation_token = kwargs.pop("continuation_token", None)
        key_name = kwargs.pop("key_name", None)
        use_managed_identity = kwargs.pop("use_managed_identity", False)
        # `sas_token` was formerly a required positional parameter
        try:
            sas_token: Optional[str] = args[0]
        except IndexError:
            sas_token = kwargs.pop("sas_token", None)

        status_response = None
        if continuation_token:
            status_response = self._use_continuation_token(continuation_token, self._client.restore_status)

        container_url, folder_name = parse_folder_url(folder_url)
        sas_parameter = self._models.SASTokenParameter(
            storage_resource_uri=container_url, token=sas_token, use_managed_identity=use_managed_identity
        )
        polling = KeyVaultBackupClientPollingMethod(
            lro_algorithms=[KeyVaultBackupClientPolling()], timeout=polling_interval, **kwargs
        )

        if key_name:
            client_method = functools.partial(self._client.begin_selective_key_restore_operation, key_name=key_name)
            restore_details = self._models.SelectiveKeyRestoreOperationParameters(
                sas_token_parameters=sas_parameter, folder=folder_name
            )
        else:
            client_method = self._client.begin_full_restore_operation
            restore_details = self._models.RestoreOperationParameters(
                sas_token_parameters=sas_parameter, folder_to_restore=folder_name
            )

        return client_method(
            restore_blob_details=restore_details,
            cls=lambda *_: None,  # poller.result() returns None
            continuation_token=status_response,
            polling=polling,
            **kwargs,
        )

    @overload
    def begin_pre_backup(
        self,
        blob_storage_url: str,
        *,
        use_managed_identity: Literal[True],
        continuation_token: Optional[str] = None,
        **kwargs: Any,
    ) -> LROPoller[None]: ...

    @overload
    def begin_pre_backup(
        self,
        blob_storage_url: str,
        *,
        sas_token: str,
        continuation_token: Optional[str] = None,
        **kwargs: Any,
    ) -> LROPoller[None]: ...

    @distributed_trace
    def begin_pre_backup(  # pylint: disable=docstring-keyword-should-match-keyword-only
        self, blob_storage_url: str, **kwargs: Any
    ) -> LROPoller[None]:
        """Initiates a pre-backup check of whether a full Key Vault backup can be performed.

        If the pre-backup check fails, calling `wait()` on the returned poller will raise an error. Otherwise, a full
        backup can be performed.

        :param str blob_storage_url: URL of the blob storage container in which the backup will be stored, for example
            https://<account>.blob.core.windows.net/backup.

        :keyword str sas_token: Optional Shared Access Signature (SAS) token to authorize access to the blob. Required
            unless `use_managed_identity` is set to True.
        :keyword use_managed_identity: Indicates which authentication method should be used. If set to True, Managed HSM
            will use the configured user-assigned managed identity to authenticate with Azure Storage. Otherwise, a SAS
            token has to be specified.
        :paramtype use_managed_identity: bool
        :keyword str continuation_token: A continuation token to restart polling from a saved state.

        :returns: An :class:`~azure.core.polling.LROPoller` instance. Call `wait()` on this object to wait for the
            operation to complete. If the check fails, an error will be raised.
        :rtype: ~azure.core.polling.LROPoller[None]
        """
        polling_interval: int = kwargs.pop("_polling_interval", 5)
        continuation_token: Optional[str] = kwargs.pop("continuation_token", None)
        use_managed_identity: bool = kwargs.pop("use_managed_identity", False)
        sas_token: Optional[str] = kwargs.pop("sas_token", None)

        parameters: PreBackupOperationParameters = PreBackupOperationParameters(
            storage_resource_uri=blob_storage_url, token=sas_token, use_managed_identity=use_managed_identity
        )
        status_response: Optional[str] = None
        if continuation_token:
            status_response = self._use_continuation_token(continuation_token, self._client.full_backup_status)

        return self._client.begin_pre_full_backup(
            pre_backup_operation_parameters=parameters,
            cls=lambda *_: None,  # poller.result() returns None
            polling=KeyVaultBackupClientPollingMethod(
                lro_algorithms=[KeyVaultBackupClientPolling()], timeout=polling_interval, **kwargs
            ),
            continuation_token=status_response,
            **kwargs,
        )

    @overload
    def begin_pre_restore(
        self,
        folder_url: str,
        *,
        use_managed_identity: Literal[True],
        continuation_token: Optional[str] = None,
        **kwargs: Any,
    ) -> LROPoller[None]: ...

    @overload
    def begin_pre_restore(
        self,
        folder_url: str,
        *,
        sas_token: str,
        continuation_token: Optional[str] = None,
        **kwargs: Any,
    ) -> LROPoller[None]: ...

    @distributed_trace
    def begin_pre_restore(  # pylint: disable=docstring-keyword-should-match-keyword-only
        self, folder_url: str, **kwargs: Any
    ) -> LROPoller[None]:
        """Initiates a pre-restore check of whether a full Key Vault restore can be performed.

        If the pre-restore check fails, calling `wait()` on the returned poller will raise an error. Otherwise, a full
        restore can be performed.

        :param str folder_url: URL of the blob holding the backup. This would be the `folder_url` of a
            :class:`KeyVaultBackupResult` returned by :func:`begin_backup`, for example
            https://<account>.blob.core.windows.net/backup/mhsm-account-2020090117323313

        :keyword str sas_token: Optional Shared Access Signature (SAS) token to authorize access to the blob. Required
            unless `use_managed_identity` is set to True.
        :keyword use_managed_identity: Indicates which authentication method should be used. If set to True, Managed HSM
            will use the configured user-assigned managed identity to authenticate with Azure Storage. Otherwise, a SAS
            token has to be specified.
        :paramtype use_managed_identity: bool
        :keyword str continuation_token: A continuation token to restart polling from a saved state.

        :returns: An :class:`~azure.core.polling.LROPoller` instance. Call `wait()` on this object to wait for the
            operation to complete. If the check fails, an error will be raised.
        :rtype: ~azure.core.polling.LROPoller[None]
        """
        polling_interval: int = kwargs.pop("_polling_interval", 5)
        continuation_token: Optional[str] = kwargs.pop("continuation_token", None)
        use_managed_identity: bool = kwargs.pop("use_managed_identity", False)
        sas_token: Optional[str] = kwargs.pop("sas_token", None)

        container_url, folder_name = parse_folder_url(folder_url)
        sas_parameter: SASTokenParameter = SASTokenParameter(
            storage_resource_uri=container_url, token=sas_token, use_managed_identity=use_managed_identity
        )
        parameters: PreRestoreOperationParameters = PreRestoreOperationParameters(
            folder_to_restore=folder_name, sas_token_parameters=sas_parameter
        )
        status_response: Optional[str] = None
        if continuation_token:
            status_response = self._use_continuation_token(continuation_token, self._client.restore_status)

        return self._client.begin_pre_full_restore_operation(
            pre_restore_operation_parameters=parameters,
            cls=lambda *_: None,  # poller.result() returns None
            polling=KeyVaultBackupClientPollingMethod(
                lro_algorithms=[KeyVaultBackupClientPolling()], timeout=polling_interval, **kwargs
            ),
            continuation_token=status_response,
            **kwargs,
        )

    def __enter__(self) -> "KeyVaultBackupClient":
        self._client.__enter__()
        return self
