# coding=utf-8
# --------------------------------------------------------------------------
# Copyright (c) Microsoft Corporation. All rights reserved.
# Licensed under the MIT License. See License.txt in the project root for license information.
# --------------------------------------------------------------------------
"""Customize generated code here.

Follow our quickstart for examples: https://aka.ms/azsdk/python/dpcodegen/python/customize
"""

import time
from typing import Any, Callable, IO, List, Optional, Union, cast

from azure.core.exceptions import ResourceNotFoundError
from azure.core.polling import PollingMethod, LROPoller, NoPolling

from azure.confidentialledger._operations._operations import (
    _ConfidentialLedgerClientOperationsMixin as GeneratedOperationsMixin,
)
from azure.confidentialledger._operations._operations import ClsType, JSON
import azure.confidentialledger.models as _models

__all__: List[str] = [
    "_ConfidentialLedgerClientOperationsMixin"
]  # Add all objects you want publicly available to users at this package level


def patch_sdk():
    """Do not remove from this file.

    `patch_sdk` is a last resort escape hatch that allows you to do customizations
    you can't accomplish using the techniques described in
    https://aka.ms/azsdk/python/dpcodegen/python/customize
    """


class BaseStatePollingMethod:
    """Base polling method for methods returning responses containing a 'state' field; the polling
    completes when 'state' becomes a desired value.
    """

    def __init__(
        self,
        operation: Callable[[], Any],
        desired_state: str,
        polling_interval_s: float,
        retry_not_found: bool,
    ):
        self._operation = operation
        self._desired_state = desired_state
        self._polling_interval_s = polling_interval_s

        self._deserialization_callback = None
        self._status = "constructed"
        self._latest_response: JSON = {}

        self._retry_not_found = retry_not_found
        self._not_found_count = 0

    def initialize(self, client, initial_response, deserialization_callback):  # pylint: disable=unused-argument
        self._evaluate_response(initial_response)
        self._deserialization_callback = deserialization_callback

    def _evaluate_response(self, response: JSON) -> None:
        self._status = "finished" if response.get("state", None) == self._desired_state else "polling"
        self._latest_response = response

    def _give_up_not_found_error(self, exception: ResourceNotFoundError) -> bool:
        if exception.error is not None and exception.error.code == "InvalidTransactionId":
            return True
        return False

    def status(self) -> str:
        return self._status

    def finished(self) -> bool:
        return self.status() in {"finished", "failed"}

    def resource(self):
        if self._deserialization_callback:
            return self._deserialization_callback(self._latest_response)
        return self._latest_response


class StatePollingMethod(BaseStatePollingMethod, PollingMethod):
    """Polling method for methods returning responses containing a 'state' field; the polling
    completes when 'state' becomes a desired value.
    """

    def __init__(
        self,
        operation: Callable[[], JSON],
        desired_state: str,
        polling_interval_s: float,
        retry_not_found: bool,
    ):
        super().__init__(operation, desired_state, polling_interval_s, retry_not_found)

    def run(self) -> None:
        try:
            while not self.finished():
                try:
                    response = self._operation()
                    self._evaluate_response(response)
                except ResourceNotFoundError as not_found_exception:
                    # We'll allow some instances of resource not found to account for replication
                    # delay if session stickiness is lost.

                    self._not_found_count += 1

                    not_retryable = not self._retry_not_found or self._give_up_not_found_error(not_found_exception)

                    if not_retryable or self._not_found_count >= 3:
                        raise
                if not self.finished():
                    time.sleep(self._polling_interval_s)
        except Exception:
            self._status = "failed"
            raise


class _ConfidentialLedgerClientOperationsMixin(GeneratedOperationsMixin):
    def begin_get_ledger_entry(
        self, transaction_id: str, *, collection_id: Optional[str] = None, **kwargs: Any
    ) -> LROPoller[_models.LedgerQueryResult]:
        """Returns a poller to fetch the ledger entry at the specified transaction id.

         A collection id may optionally be specified to indicate the collection from which to fetch
         the value.

        To return older ledger entries, the relevant sections of the ledger must be
         read from disk and validated. To prevent blocking within the enclave, the
         response will indicate whether the entry is ready and part of the response, or
         if the loading is still ongoing.

         :param transaction_id: Identifies a write transaction. Required.
         :type transaction_id: str
         :keyword collection_id: The collection id. Default value is None.
         :paramtype collection_id: str
         :return: An instance of LROPoller that returns a LedgerQueryResult object for the ledger entry.
         :rtype: ~azure.core.polling.LROPoller[~azure.confidentialledger.models.LedgerQueryResult]
         :raises ~azure.core.exceptions.HttpResponseError:
        """
        polling = kwargs.pop("polling", True)  # type: Union[bool, PollingMethod]
        lro_delay = kwargs.pop("polling_interval", 0.5)

        def operation() -> JSON:
            return super(_ConfidentialLedgerClientOperationsMixin, self).get_ledger_entry(
                transaction_id, collection_id=collection_id, **kwargs
            )

        initial_response = operation()

        if polling is True:
            polling_method = cast(PollingMethod, StatePollingMethod(operation, "Ready", lro_delay, False))
        elif polling is False:
            polling_method = cast(PollingMethod, NoPolling())
        else:
            polling_method = polling
        return LROPoller(self._client, initial_response, lambda x: x, polling_method)

    def begin_get_receipt(self, transaction_id: str, **kwargs: Any) -> LROPoller[_models.TransactionReceipt]:
        """Returns a poller for getting a receipt certifying ledger contents at a particular
        transaction id.

        :param transaction_id: Identifies a write transaction. Required.
        :type transaction_id: str
        :return: An instance of LROPoller that returns a TransactionReceipt object for the receipt.
        :rtype: ~azure.core.polling.LROPoller[~azure.confidentialledger.models.TransactionReceipt]
        :raises ~azure.core.exceptions.HttpResponseError:
        """
        polling = kwargs.pop("polling", True)  # type: Union[bool, PollingMethod]
        lro_delay = kwargs.pop("polling_interval", 0.5)

        def operation() -> JSON:
            return super(_ConfidentialLedgerClientOperationsMixin, self).get_receipt(
                transaction_id=transaction_id, **kwargs
            )

        initial_response = operation()

        if polling is True:
            polling_method = cast(PollingMethod, StatePollingMethod(operation, "Ready", lro_delay, False))
        elif polling is False:
            polling_method = cast(PollingMethod, NoPolling())
        else:
            polling_method = polling
        return LROPoller(self._client, initial_response, lambda x: x, polling_method)

    def begin_create_ledger_entry(
        self,
        entry: Union[_models.LedgerEntry, JSON, IO[bytes]],
        *,
        collection_id: Optional[str] = None,
        **kwargs: Any,
    ) -> LROPoller[_models.TransactionStatus]:
        """Writes a ledger entry and returns a poller to wait for it to be durably committed. The
        poller returns the result for the initial call to create the ledger entry.

        A collection id may optionally be specified.

        :param entry: Ledger entry. Is one of the following types: LedgerEntry, JSON, IO[bytes]
         Required.
        :type entry: ~azure.confidentialledger.models.LedgerEntry or JSON or IO[bytes]
        :keyword collection_id: The collection id. Default value is None.
        :paramtype collection_id: str
        :return: TransactionStatus. The TransactionStatus is compatible with MutableMapping
        :rtype: ~azure.core.polling.LROPoller[~azure.confidentialledger.models.TransactionStatus]
        :raises ~azure.core.exceptions.HttpResponseError:
        """

        # Pop arguments that are unexpected in the pipeline.

        polling = kwargs.pop("polling", True)  # type: Union[bool, PollingMethod]
        lro_delay = kwargs.pop("polling_interval", 0.5)

        # Pop the custom deserializer, if any, so we know the format of the response and can
        # retrieve the transactionId. Serialize the response later.

        cls = kwargs.pop("cls", None)  # type: ClsType[JSON]
        kwargs["cls"] = lambda pipeline_response, json_response, headers: (
            pipeline_response,
            {
                **json_response,
                "transactionId": headers.get("x-ms-ccf-transaction-id") if headers else None,
            },
            headers,
        )

        post_pipeline_response, post_result, post_headers = self.create_ledger_entry(
            entry, collection_id=collection_id, **kwargs
        )

        # Delete the cls because it should only apply to the create_ledger_entry response, not the
        # wait_for_commit call.

        del kwargs["cls"]

        transaction_id = post_result["transactionId"]  # type: ignore

        kwargs["polling"] = polling
        kwargs["polling_interval"] = lro_delay

        if cls:
            kwargs["_create_ledger_entry_response"] = cls(
                post_pipeline_response, cast(JSON, post_result), post_headers  # type: ignore
            )
        else:
            kwargs["_create_ledger_entry_response"] = post_result
        return self.begin_wait_for_commit(transaction_id, **kwargs)

    def begin_wait_for_commit(
        self,
        transaction_id,  # type: str
        **kwargs,  # type: Any
    ) -> LROPoller[_models.TransactionStatus]:
        """Creates a poller that queries the state of the specified transaction until it is
        Committed, a state that indicates the transaction is durably stored in the Confidential
        Ledger.

        :param transaction_id: Identifies a write transaction. Required.
        :type transaction_id: str
        :return: An instance of LROPoller returning a TransactionStatus object describing the transaction status.
        :rtype: ~azure.core.polling.LROPoller[~azure.confidentialledger.models.TransactionStatus]
        :raises ~azure.core.exceptions.HttpResponseError:
        """
        polling = kwargs.pop("polling", True)  # type: Union[bool, PollingMethod]
        lro_delay = kwargs.pop("polling_interval", 0.5)

        # If this poller was called from begin_create_ledger_entry, we should return the
        # create_ledger_entry response, not the transaction status.

        post_result = kwargs.pop("_create_ledger_entry_response", None)

        def deserialization_callback(x):
            return x if post_result is None else post_result

        def operation() -> JSON:
            return super(_ConfidentialLedgerClientOperationsMixin, self).get_transaction_status(
                transaction_id=transaction_id, **kwargs
            )

        try:
            initial_response = operation()
        except ResourceNotFoundError:
            if polling is False or polling is None:
                raise
            # This method allows for temporary resource not found errors, which may occur if session
            # stickiness is lost and there is replication lag.

            initial_response = {}
        if polling is True:
            polling_method = cast(PollingMethod, StatePollingMethod(operation, "Committed", lro_delay, True))
        elif polling is False:
            polling_method = cast(PollingMethod, NoPolling())
        else:
            polling_method = polling
        return LROPoller(self._client, initial_response, deserialization_callback, polling_method)

    def create_ledger_entry(
        self,
        entry: Union[_models.LedgerEntry, JSON, IO[bytes]],
        *,
        collection_id: Optional[str] = None,
        **kwargs: Any,
    ) -> _models.LedgerWriteResult:
        """Writes a ledger entry.

        The result is the expected JSON response with an additional field
        'transactionId' which represents the transaction identifier for this write operation.

        A collection id may optionally be specified.

        :param entry: Ledger entry. Is one of the following types: LedgerEntry, JSON, IO[bytes]
         Required.
        :type entry: ~azure.confidentialledger.models.LedgerEntry or JSON or IO[bytes]
        :keyword collection_id: The collection id. Default value is None.
        :paramtype collection_id: str
        :return: LedgerWriteResult. The LedgerWriteResult is compatible with MutableMapping
        :rtype: ~azure.confidentialledger.models.LedgerWriteResult
        :raises ~azure.core.exceptions.HttpResponseError:

        Example:
            .. code-block:: python

                # JSON input template you can fill out and use as your body input.
                entry = {
                    "collectionId": {
                        "collectionId": "str"  # Required.
                    },
                    "contents": "str",  # Required. Contents of the ledger entry.
                    "transactionId": "str"  # Optional. A unique identifier for the state of the
                      ledger. If returned as part of a LedgerEntry, it indicates the state from which
                      the entry was read.
                }
        """

        kwargs["cls"] = kwargs.get(
            "cls",
            lambda _, json_response, headers: {
                **json_response,
                "transactionId": headers.get("x-ms-ccf-transaction-id") if headers else None,
            },
        )
        return super().create_ledger_entry(entry, collection_id=collection_id, **kwargs)
