# pylint: disable=line-too-long,useless-suppression
# ------------------------------------
# Copyright (c) Microsoft Corporation.
# Licensed under the MIT License.
# ------------------------------------
"""Customize generated code here.

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

from typing import Final, FrozenSet, List, Dict, Mapping, Optional, Any, Tuple
from azure.core.polling import LROPoller, AsyncLROPoller, PollingMethod, AsyncPollingMethod
from azure.core.polling.base_polling import (
    LROBasePolling,
    OperationFailed,
    _raise_if_bad_http_status_and_method,
)
from azure.core.polling.async_base_polling import AsyncLROBasePolling
from ._patch_evaluation_typeddicts import (
    AzureAIAgentTargetParam,
    AzureAIBenchmarkPreviewEvalRunDataSource,
    AzureAIDataSourceConfig,
    AzureAIModelTargetParam,
    AzureAIResponsesEvalRunDataSource,
    EvalCsvFileIdSource,
    EvalCsvRunDataSource,
    TestingCriterionAzureAIEvaluator,
    ModelSamplingConfigParam,
    RedTeamEvalRunDataSource,
    ResponseRetrievalItemGenerationParams,
    TargetCompletionEvalRunDataSource,
    ToolDescriptionParam,
    TracesPreviewEvalRunDataSource,
)
from ._models import CustomCredential as CustomCredentialGenerated
from ..models import MemoryStoreUpdateCompletedResult, MemoryStoreUpdateResult
from ._enums import _FoundryFeaturesOptInKeys, _AgentDefinitionOptInKeys

_FOUNDRY_FEATURES_HEADER_NAME: Final[str] = "Foundry-Features"
"""The HTTP header name used to opt in to Foundry preview features."""

_BETA_OPERATION_FEATURE_HEADERS: Final[dict] = {
    "evaluation_taxonomies": _FoundryFeaturesOptInKeys.EVALUATIONS_V1_PREVIEW.value,
    "evaluators": _FoundryFeaturesOptInKeys.EVALUATIONS_V1_PREVIEW.value,
    "insights": _FoundryFeaturesOptInKeys.INSIGHTS_V1_PREVIEW.value,
    "memory_stores": _FoundryFeaturesOptInKeys.MEMORY_STORES_V1_PREVIEW.value,
    "red_teams": _FoundryFeaturesOptInKeys.RED_TEAMS_V1_PREVIEW.value,
    "schedules": _FoundryFeaturesOptInKeys.SCHEDULES_V1_PREVIEW.value,
    "toolboxes": _FoundryFeaturesOptInKeys.TOOLBOXES_V1_PREVIEW.value,
    "skills": _FoundryFeaturesOptInKeys.SKILLS_V1_PREVIEW.value,
    "agents": ",".join(
        [
            _AgentDefinitionOptInKeys.HOSTED_AGENTS_V1_PREVIEW.value,
            _AgentDefinitionOptInKeys.AGENT_ENDPOINT_V1_PREVIEW.value,
        ]
    ),
}
"""Foundry-Features header values keyed by beta sub-client property name."""


def _has_header_case_insensitive(headers: Any, header_name: str) -> bool:
    """Return True if headers already contains the provided header name.

    :param headers: The headers mapping to search.
    :type headers: Any
    :param header_name: The header name to look for (case-insensitive).
    :type header_name: str
    :return: True if the header is present, False otherwise.
    :rtype: bool
    """
    try:
        header_name_lower = header_name.lower()
        return any(str(key).lower() == header_name_lower for key in headers)
    except Exception:  # pylint: disable=broad-except
        return False


class CustomCredential(CustomCredentialGenerated, discriminator="CustomKeys"):
    """Custom credential definition.

    :ivar type: The credential type. Always equals CredentialType.CUSTOM. Required.
    :vartype type: str or ~azure.ai.projects.models.CredentialType
    :ivar credential_keys: The secret custom credential keys. Required.
    :vartype credential_keys: dict[str, str]
    """

    credential_keys: Dict[str, str]
    """The secret custom credential keys. Required."""

    def __init__(self, *args: Any, **kwargs: Any) -> None:
        super().__init__(*args, **kwargs)

        # Fix for GitHub issue https://github.com/Azure/azure-sdk-for-net/issues/52355
        # Although the issue was filed on C# Projects SDK, the same problem exists in Python SDK.
        # Assume your Foundry project has a connection of type `Custom`, named "test_custom_connection",
        # and you defined two public and two secrete (private) keys. When you get the connection, the response
        # payload will look something like this:
        #     {
        #         "name": "test_custom_connection",
        #         "id": "/subscriptions/.../connections/test_custom_connection",
        #         "type": "CustomKeys",
        #         "target": "_",
        #         "isDefault": true,
        #         "credentials": {
        #             "nameofprivatekey1": "PrivateKey1",
        #             "nameofprivatekey2": "PrivateKey2",
        #             "type": "CustomKeys"
        #         },
        #         "metadata": {
        #             "NameOfPublicKey1": "PublicKey1",
        #             "NameOfPublicKey2": "PublicKey2"
        #         }
        #     }
        # We would like to add a new Dict property on the Python `credentials` object, named `credential_keys`,
        # to hold all the secret keys. This is done by the line below.
        if args and isinstance(args[0], Mapping):
            self.credential_keys = {k: v for k, v in args[0].items() if k != "type"}
        else:
            self.credential_keys = {}


_FINISHED: Final[FrozenSet[str]] = frozenset(["completed", "superseded", "failed"])
_FAILED: Final[FrozenSet[str]] = frozenset(["failed"])


class _UpdateMemoriesLROPollingMethod(LROBasePolling):
    """A custom polling method implementation for Memory Store updates."""

    @property
    def _current_body(self) -> MemoryStoreUpdateResult:
        try:
            return MemoryStoreUpdateResult(self._pipeline_response.http_response.json())
        except Exception:  # pylint: disable=broad-exception-caught
            return MemoryStoreUpdateResult()  # type: ignore[call-overload]

    def finished(self) -> bool:
        """Is this polling finished?

        :return: True/False for whether polling is complete.
        :rtype: bool
        """
        return self._finished(self.status())

    @staticmethod
    def _finished(status) -> bool:
        if hasattr(status, "value"):
            status = status.value
        return str(status).lower() in _FINISHED

    @staticmethod
    def _failed(status) -> bool:
        if hasattr(status, "value"):
            status = status.value
        return str(status).lower() in _FAILED

    def get_continuation_token(self) -> str:
        return self._current_body.update_id

    # pylint: disable=arguments-differ
    def from_continuation_token(self, continuation_token: str, **kwargs: Any) -> Tuple:  # type: ignore[override]
        try:
            client = kwargs["client"]
        except KeyError as exc:
            raise ValueError("Need kwarg 'client' to be recreated from continuation_token") from exc

        try:
            deserialization_callback = kwargs["deserialization_callback"]
        except KeyError as exc:
            raise ValueError("Need kwarg 'deserialization_callback' to be recreated from continuation_token") from exc

        return client, continuation_token, deserialization_callback

    def _poll(self) -> None:
        """Poll status of operation so long as operation is incomplete and
        we have an endpoint to query.

        :raises: OperationFailed if operation status 'Failed' or 'Canceled'.
        :raises: BadStatus if response status invalid.
        :raises: BadResponse if response invalid.
        """

        if not self.finished():
            self.update_status()
        while not self.finished():
            self._delay()
            self.update_status()

        if self._failed(self.status()):
            raise OperationFailed("Operation failed or canceled")

        final_get_url = self._operation.get_final_get_url(self._pipeline_response)
        if final_get_url:
            self._pipeline_response = self.request_status(final_get_url)
            _raise_if_bad_http_status_and_method(self._pipeline_response.http_response)


class _AsyncUpdateMemoriesLROPollingMethod(AsyncLROBasePolling):
    """A custom polling method implementation for Memory Store updates."""

    @property
    def _current_body(self) -> MemoryStoreUpdateResult:
        try:
            return MemoryStoreUpdateResult(self._pipeline_response.http_response.json())
        except Exception:  # pylint: disable=broad-exception-caught
            return MemoryStoreUpdateResult()  # type: ignore[call-overload]

    def finished(self) -> bool:
        """Is this polling finished?

        :return: True/False for whether polling is complete.
        :rtype: bool
        """
        return self._finished(self.status())

    @staticmethod
    def _finished(status) -> bool:
        if hasattr(status, "value"):
            status = status.value
        return str(status).lower() in _FINISHED

    @staticmethod
    def _failed(status) -> bool:
        if hasattr(status, "value"):
            status = status.value
        return str(status).lower() in _FAILED

    def get_continuation_token(self) -> str:
        return self._current_body.update_id

    # pylint: disable=arguments-differ
    def from_continuation_token(self, continuation_token: str, **kwargs: Any) -> Tuple:  # type: ignore[override]
        try:
            client = kwargs["client"]
        except KeyError as exc:
            raise ValueError("Need kwarg 'client' to be recreated from continuation_token") from exc

        try:
            deserialization_callback = kwargs["deserialization_callback"]
        except KeyError as exc:
            raise ValueError("Need kwarg 'deserialization_callback' to be recreated from continuation_token") from exc

        return client, continuation_token, deserialization_callback

    async def _poll(self) -> None:
        """Poll status of operation so long as operation is incomplete and
        we have an endpoint to query.

        :raises: OperationFailed if operation status 'Failed' or 'Canceled'.
        :raises: BadStatus if response status invalid.
        :raises: BadResponse if response invalid.
        """

        if not self.finished():
            await self.update_status()
        while not self.finished():
            await self._delay()
            await self.update_status()

        if self._failed(self.status()):
            raise OperationFailed("Operation failed or canceled")

        final_get_url = self._operation.get_final_get_url(self._pipeline_response)
        if final_get_url:
            self._pipeline_response = await self.request_status(final_get_url)
            _raise_if_bad_http_status_and_method(self._pipeline_response.http_response)


class UpdateMemoriesLROPoller(LROPoller[MemoryStoreUpdateCompletedResult]):
    """Custom LROPoller for Memory Store update operations."""

    _polling_method: "_UpdateMemoriesLROPollingMethod"

    @property
    def update_id(self) -> str:
        """Returns the update ID associated with the long-running update memories operation.

        :return: Returns the update ID.
        :rtype: str
        """
        return self._polling_method._current_body.update_id  # pylint: disable=protected-access

    @property
    def superseded_by(self) -> Optional[str]:
        """Returns the ID of the operation that superseded this update.

        :return: Returns the ID of the superseding operation, if it exists.
        :rtype: Optional[str]
        """
        return (
            self._polling_method._current_body.superseded_by  # pylint: disable=protected-access
            if self._polling_method._current_body  # pylint: disable=protected-access
            else None
        )

    @classmethod
    def from_continuation_token(
        cls, polling_method: PollingMethod[MemoryStoreUpdateCompletedResult], continuation_token: str, **kwargs: Any
    ) -> "UpdateMemoriesLROPoller":
        """Create a poller from a continuation token.

        :param polling_method: The polling strategy to adopt
        :type polling_method: ~azure.core.polling.PollingMethod
        :param continuation_token: An opaque continuation token
        :type continuation_token: str
        :return: An instance of UpdateMemoriesLROPoller
        :rtype: UpdateMemoriesLROPoller
        :raises ~azure.core.exceptions.HttpResponseError: If the continuation token is invalid.
        """
        (
            client,
            initial_response,
            deserialization_callback,
        ) = polling_method.from_continuation_token(continuation_token, **kwargs)

        return cls(client, initial_response, deserialization_callback, polling_method)


class AsyncUpdateMemoriesLROPoller(AsyncLROPoller[MemoryStoreUpdateCompletedResult]):
    """Custom AsyncLROPoller for Memory Store update operations."""

    _polling_method: "_AsyncUpdateMemoriesLROPollingMethod"

    @property
    def update_id(self) -> str:
        """Returns the update ID associated with the long-running update memories operation.

        :return: Returns the update ID.
        :rtype: str
        """
        return self._polling_method._current_body.update_id  # pylint: disable=protected-access

    @property
    def superseded_by(self) -> Optional[str]:
        """Returns the ID of the operation that superseded this update.

        :return: Returns the ID of the superseding operation, if it exists.
        :rtype: Optional[str]
        """
        return (
            self._polling_method._current_body.superseded_by  # pylint: disable=protected-access
            if self._polling_method._current_body  # pylint: disable=protected-access
            else None
        )

    @classmethod
    def from_continuation_token(
        cls,
        polling_method: AsyncPollingMethod[MemoryStoreUpdateCompletedResult],
        continuation_token: str,
        **kwargs: Any,
    ) -> "AsyncUpdateMemoriesLROPoller":
        """Create a poller from a continuation token.

        :param polling_method: The polling strategy to adopt
        :type polling_method: ~azure.core.polling.PollingMethod
        :param continuation_token: An opaque continuation token
        :type continuation_token: str
        :return: An instance of AsyncUpdateMemoriesLROPoller
        :rtype: AsyncUpdateMemoriesLROPoller
        :raises ~azure.core.exceptions.HttpResponseError: If the continuation token is invalid.
        """
        (
            client,
            initial_response,
            deserialization_callback,
        ) = polling_method.from_continuation_token(continuation_token, **kwargs)

        return cls(client, initial_response, deserialization_callback, polling_method)


__all__: List[str] = [
    "AsyncUpdateMemoriesLROPoller",
    "AzureAIAgentTargetParam",
    "AzureAIBenchmarkPreviewEvalRunDataSource",
    "AzureAIDataSourceConfig",
    "AzureAIModelTargetParam",
    "AzureAIResponsesEvalRunDataSource",
    "CustomCredential",
    "EvalCsvFileIdSource",
    "EvalCsvRunDataSource",
    "TestingCriterionAzureAIEvaluator",
    "ModelSamplingConfigParam",
    "RedTeamEvalRunDataSource",
    "ResponseRetrievalItemGenerationParams",
    "TargetCompletionEvalRunDataSource",
    "ToolDescriptionParam",
    "TracesPreviewEvalRunDataSource",
    "UpdateMemoriesLROPoller",
]  # 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
    """
