# --------------------------------------------------------------------------------------------
# Copyright (c) Microsoft Corporation. All rights reserved.
# Licensed under the MIT License. See License.txt in the project root for license information.
#
# Code generated by aaz-dev-tools
# --------------------------------------------------------------------------------------------

# pylint: skip-file
# flake8: noqa

from azure.cli.core.aaz import *


@register_command(
    "disk-encryption-set update",
)
class Update(AAZCommand):
    """Update a disk encryption set

    :example: Update a disk encryption set.
        az disk-encryption-set update --name MyDiskEncryptionSet --resource-group MyResourceGroup --key-url MyKey --source-vault MyVault

    :example: Update multi-tenant application client id of a disk encryption set.
        az disk-encryption-set update --name MyDiskEncryptionSet --resource-group MyResourceGroup --key-url MyKey --source-vault MyVault --federated-client-id myFederatedClientId

    :example: Clear multi-tenant application client id of a disk encryption set.
        az disk-encryption-set update --name MyDiskEncryptionSet --resource-group MyResourceGroup --key-url MyKey --source-vault MyVault --federated-client-id None
    """

    _aaz_info = {
        "version": "2023-04-02",
        "resources": [
            ["mgmt-plane", "/subscriptions/{}/resourcegroups/{}/providers/microsoft.compute/diskencryptionsets/{}", "2023-04-02"],
        ]
    }

    AZ_SUPPORT_NO_WAIT = True

    AZ_SUPPORT_GENERIC_UPDATE = True

    def _handler(self, command_args):
        super()._handler(command_args)
        return self.build_lro_poller(self._execute_operations, self._output)

    _args_schema = None

    @classmethod
    def _build_arguments_schema(cls, *args, **kwargs):
        if cls._args_schema is not None:
            return cls._args_schema
        cls._args_schema = super()._build_arguments_schema(*args, **kwargs)

        # define Arg Group ""

        _args_schema = cls._args_schema
        _args_schema.disk_encryption_set_name = AAZStrArg(
            options=["-n", "--name", "--disk-encryption-set-name"],
            help="Name of disk encryption set.",
            required=True,
            id_part="name",
        )
        _args_schema.resource_group = AAZResourceGroupNameArg(
            required=True,
        )
        _args_schema.key_url = AAZStrArg(
            options=["--key-url"],
            help="URL pointing to a key or secret in KeyVault.",
        )
        _args_schema.source_vault = AAZStrArg(
            options=["--source-vault"],
            help="Name or ID of the KeyVault containing the key or secret.",
            nullable=True,
        )
        _args_schema.federated_client_id = AAZStrArg(
            options=["--federated-client-id"],
            help="The federated client id used in cross tenant scenario.",
            nullable=True,
        )
        _args_schema.enable_auto_key_rotation = AAZBoolArg(
            options=["--auto-rotation", "--enable-auto-key-rotation"],
            help="Enable automatic rotation of keys.",
            nullable=True,
        )

        # define Arg Group "DiskEncryptionSet"

        _args_schema = cls._args_schema
        _args_schema.identity = AAZObjectArg(
            options=["--identity"],
            arg_group="DiskEncryptionSet",
            help="The managed identity for the disk encryption set. It should be given permission on the key vault before it can be used to encrypt disks.",
            nullable=True,
        )

        identity = cls._args_schema.identity
        identity.type = AAZStrArg(
            options=["type"],
            help="The type of Managed Identity used by the DiskEncryptionSet. Only SystemAssigned is supported for new creations. Disk Encryption Sets can be updated with Identity type None during migration of subscription to a new Azure Active Directory tenant; it will cause the encrypted resources to lose access to the keys.",
            nullable=True,
            enum={"None": "None", "SystemAssigned": "SystemAssigned", "SystemAssigned, UserAssigned": "SystemAssigned, UserAssigned", "UserAssigned": "UserAssigned"},
        )
        identity.user_assigned_identities = AAZDictArg(
            options=["user-assigned-identities"],
            help="The list of user identities associated with the disk encryption set. The user identity dictionary key references will be ARM resource ids in the form: '/subscriptions/{subscriptionId}/resourceGroups/{resourceGroupName}/providers/Microsoft.ManagedIdentity/userAssignedIdentities/{identityName}'.",
            nullable=True,
        )

        user_assigned_identities = cls._args_schema.identity.user_assigned_identities
        user_assigned_identities.Element = AAZObjectArg(
            nullable=True,
            blank={},
        )
        return cls._args_schema

    def _execute_operations(self):
        self.pre_operations()
        self.DiskEncryptionSetsGet(ctx=self.ctx)()
        self.pre_instance_update(self.ctx.vars.instance)
        self.InstanceUpdateByJson(ctx=self.ctx)()
        self.InstanceUpdateByGeneric(ctx=self.ctx)()
        self.post_instance_update(self.ctx.vars.instance)
        yield self.DiskEncryptionSetsCreateOrUpdate(ctx=self.ctx)()
        self.post_operations()

    @register_callback
    def pre_operations(self):
        pass

    @register_callback
    def post_operations(self):
        pass

    @register_callback
    def pre_instance_update(self, instance):
        pass

    @register_callback
    def post_instance_update(self, instance):
        pass

    def _output(self, *args, **kwargs):
        result = self.deserialize_output(self.ctx.vars.instance, client_flatten=True)
        return result

    class DiskEncryptionSetsGet(AAZHttpOperation):
        CLIENT_TYPE = "MgmtClient"

        def __call__(self, *args, **kwargs):
            request = self.make_request()
            session = self.client.send_request(request=request, stream=False, **kwargs)
            if session.http_response.status_code in [200]:
                return self.on_200(session)

            return self.on_error(session.http_response)

        @property
        def url(self):
            return self.client.format_url(
                "/subscriptions/{subscriptionId}/resourceGroups/{resourceGroupName}/providers/Microsoft.Compute/diskEncryptionSets/{diskEncryptionSetName}",
                **self.url_parameters
            )

        @property
        def method(self):
            return "GET"

        @property
        def error_format(self):
            return "ODataV4Format"

        @property
        def url_parameters(self):
            parameters = {
                **self.serialize_url_param(
                    "diskEncryptionSetName", self.ctx.args.disk_encryption_set_name,
                    required=True,
                ),
                **self.serialize_url_param(
                    "resourceGroupName", self.ctx.args.resource_group,
                    required=True,
                ),
                **self.serialize_url_param(
                    "subscriptionId", self.ctx.subscription_id,
                    required=True,
                ),
            }
            return parameters

        @property
        def query_parameters(self):
            parameters = {
                **self.serialize_query_param(
                    "api-version", "2023-04-02",
                    required=True,
                ),
            }
            return parameters

        @property
        def header_parameters(self):
            parameters = {
                **self.serialize_header_param(
                    "Accept", "application/json",
                ),
            }
            return parameters

        def on_200(self, session):
            data = self.deserialize_http_content(session)
            self.ctx.set_var(
                "instance",
                data,
                schema_builder=self._build_schema_on_200
            )

        _schema_on_200 = None

        @classmethod
        def _build_schema_on_200(cls):
            if cls._schema_on_200 is not None:
                return cls._schema_on_200

            cls._schema_on_200 = AAZObjectType()
            _UpdateHelper._build_schema_disk_encryption_set_read(cls._schema_on_200)

            return cls._schema_on_200

    class DiskEncryptionSetsCreateOrUpdate(AAZHttpOperation):
        CLIENT_TYPE = "MgmtClient"

        def __call__(self, *args, **kwargs):
            request = self.make_request()
            session = self.client.send_request(request=request, stream=False, **kwargs)
            if session.http_response.status_code in [202]:
                return self.client.build_lro_polling(
                    self.ctx.args.no_wait,
                    session,
                    self.on_200,
                    self.on_error,
                    lro_options={"final-state-via": "azure-async-operation"},
                    path_format_arguments=self.url_parameters,
                )
            if session.http_response.status_code in [200]:
                return self.client.build_lro_polling(
                    self.ctx.args.no_wait,
                    session,
                    self.on_200,
                    self.on_error,
                    lro_options={"final-state-via": "azure-async-operation"},
                    path_format_arguments=self.url_parameters,
                )

            return self.on_error(session.http_response)

        @property
        def url(self):
            return self.client.format_url(
                "/subscriptions/{subscriptionId}/resourceGroups/{resourceGroupName}/providers/Microsoft.Compute/diskEncryptionSets/{diskEncryptionSetName}",
                **self.url_parameters
            )

        @property
        def method(self):
            return "PUT"

        @property
        def error_format(self):
            return "ODataV4Format"

        @property
        def url_parameters(self):
            parameters = {
                **self.serialize_url_param(
                    "diskEncryptionSetName", self.ctx.args.disk_encryption_set_name,
                    required=True,
                ),
                **self.serialize_url_param(
                    "resourceGroupName", self.ctx.args.resource_group,
                    required=True,
                ),
                **self.serialize_url_param(
                    "subscriptionId", self.ctx.subscription_id,
                    required=True,
                ),
            }
            return parameters

        @property
        def query_parameters(self):
            parameters = {
                **self.serialize_query_param(
                    "api-version", "2023-04-02",
                    required=True,
                ),
            }
            return parameters

        @property
        def header_parameters(self):
            parameters = {
                **self.serialize_header_param(
                    "Content-Type", "application/json",
                ),
                **self.serialize_header_param(
                    "Accept", "application/json",
                ),
            }
            return parameters

        @property
        def content(self):
            _content_value, _builder = self.new_content_builder(
                self.ctx.args,
                value=self.ctx.vars.instance,
            )

            return self.serialize_content(_content_value)

        def on_200(self, session):
            data = self.deserialize_http_content(session)
            self.ctx.set_var(
                "instance",
                data,
                schema_builder=self._build_schema_on_200
            )

        _schema_on_200 = None

        @classmethod
        def _build_schema_on_200(cls):
            if cls._schema_on_200 is not None:
                return cls._schema_on_200

            cls._schema_on_200 = AAZObjectType()
            _UpdateHelper._build_schema_disk_encryption_set_read(cls._schema_on_200)

            return cls._schema_on_200

    class InstanceUpdateByJson(AAZJsonInstanceUpdateOperation):

        def __call__(self, *args, **kwargs):
            self._update_instance(self.ctx.vars.instance)

        def _update_instance(self, instance):
            _instance_value, _builder = self.new_content_builder(
                self.ctx.args,
                value=instance,
                typ=AAZObjectType
            )
            _builder.set_prop("identity", AAZIdentityObjectType, ".identity")
            _builder.set_prop("properties", AAZObjectType, typ_kwargs={"flags": {"client_flatten": True}})

            identity = _builder.get(".identity")
            if identity is not None:
                identity.set_prop("type", AAZStrType, ".type")
                identity.set_prop("userAssignedIdentities", AAZDictType, ".user_assigned_identities")

            user_assigned_identities = _builder.get(".identity.userAssignedIdentities")
            if user_assigned_identities is not None:
                user_assigned_identities.set_elements(AAZObjectType, ".")

            properties = _builder.get(".properties")
            if properties is not None:
                properties.set_prop("activeKey", AAZObjectType)
                properties.set_prop("federatedClientId", AAZStrType, ".federated_client_id")
                properties.set_prop("rotationToLatestKeyVersionEnabled", AAZBoolType, ".enable_auto_key_rotation")

            active_key = _builder.get(".properties.activeKey")
            if active_key is not None:
                active_key.set_prop("keyUrl", AAZStrType, ".key_url", typ_kwargs={"flags": {"required": True}})
                active_key.set_prop("sourceVault", AAZObjectType)

            source_vault = _builder.get(".properties.activeKey.sourceVault")
            if source_vault is not None:
                source_vault.set_prop("id", AAZStrType, ".source_vault")

            return _instance_value

    class InstanceUpdateByGeneric(AAZGenericInstanceUpdateOperation):

        def __call__(self, *args, **kwargs):
            self._update_instance_by_generic(
                self.ctx.vars.instance,
                self.ctx.generic_update_args
            )


class _UpdateHelper:
    """Helper class for Update"""

    _schema_api_error_read = None

    @classmethod
    def _build_schema_api_error_read(cls, _schema):
        if cls._schema_api_error_read is not None:
            _schema.code = cls._schema_api_error_read.code
            _schema.details = cls._schema_api_error_read.details
            _schema.innererror = cls._schema_api_error_read.innererror
            _schema.message = cls._schema_api_error_read.message
            _schema.target = cls._schema_api_error_read.target
            return

        cls._schema_api_error_read = _schema_api_error_read = AAZObjectType(
            flags={"read_only": True}
        )

        api_error_read = _schema_api_error_read
        api_error_read.code = AAZStrType()
        api_error_read.details = AAZListType()
        api_error_read.innererror = AAZObjectType()
        api_error_read.message = AAZStrType()
        api_error_read.target = AAZStrType()

        details = _schema_api_error_read.details
        details.Element = AAZObjectType()

        _element = _schema_api_error_read.details.Element
        _element.code = AAZStrType()
        _element.message = AAZStrType()
        _element.target = AAZStrType()

        innererror = _schema_api_error_read.innererror
        innererror.errordetail = AAZStrType()
        innererror.exceptiontype = AAZStrType()

        _schema.code = cls._schema_api_error_read.code
        _schema.details = cls._schema_api_error_read.details
        _schema.innererror = cls._schema_api_error_read.innererror
        _schema.message = cls._schema_api_error_read.message
        _schema.target = cls._schema_api_error_read.target

    _schema_disk_encryption_set_read = None

    @classmethod
    def _build_schema_disk_encryption_set_read(cls, _schema):
        if cls._schema_disk_encryption_set_read is not None:
            _schema.id = cls._schema_disk_encryption_set_read.id
            _schema.identity = cls._schema_disk_encryption_set_read.identity
            _schema.location = cls._schema_disk_encryption_set_read.location
            _schema.name = cls._schema_disk_encryption_set_read.name
            _schema.properties = cls._schema_disk_encryption_set_read.properties
            _schema.tags = cls._schema_disk_encryption_set_read.tags
            _schema.type = cls._schema_disk_encryption_set_read.type
            return

        cls._schema_disk_encryption_set_read = _schema_disk_encryption_set_read = AAZObjectType()

        disk_encryption_set_read = _schema_disk_encryption_set_read
        disk_encryption_set_read.id = AAZStrType(
            flags={"read_only": True},
        )
        disk_encryption_set_read.identity = AAZIdentityObjectType()
        disk_encryption_set_read.location = AAZStrType(
            flags={"required": True},
        )
        disk_encryption_set_read.name = AAZStrType(
            flags={"read_only": True},
        )
        disk_encryption_set_read.properties = AAZObjectType(
            flags={"client_flatten": True},
        )
        disk_encryption_set_read.tags = AAZDictType()
        disk_encryption_set_read.type = AAZStrType(
            flags={"read_only": True},
        )

        identity = _schema_disk_encryption_set_read.identity
        identity.principal_id = AAZStrType(
            serialized_name="principalId",
            flags={"read_only": True},
        )
        identity.tenant_id = AAZStrType(
            serialized_name="tenantId",
            flags={"read_only": True},
        )
        identity.type = AAZStrType()
        identity.user_assigned_identities = AAZDictType(
            serialized_name="userAssignedIdentities",
        )

        user_assigned_identities = _schema_disk_encryption_set_read.identity.user_assigned_identities
        user_assigned_identities.Element = AAZObjectType()

        _element = _schema_disk_encryption_set_read.identity.user_assigned_identities.Element
        _element.client_id = AAZStrType(
            serialized_name="clientId",
            flags={"read_only": True},
        )
        _element.principal_id = AAZStrType(
            serialized_name="principalId",
            flags={"read_only": True},
        )

        properties = _schema_disk_encryption_set_read.properties
        properties.active_key = AAZObjectType(
            serialized_name="activeKey",
        )
        cls._build_schema_key_for_disk_encryption_set_read(properties.active_key)
        properties.auto_key_rotation_error = AAZObjectType(
            serialized_name="autoKeyRotationError",
            flags={"read_only": True},
        )
        cls._build_schema_api_error_read(properties.auto_key_rotation_error)
        properties.encryption_type = AAZStrType(
            serialized_name="encryptionType",
        )
        properties.federated_client_id = AAZStrType(
            serialized_name="federatedClientId",
        )
        properties.last_key_rotation_timestamp = AAZStrType(
            serialized_name="lastKeyRotationTimestamp",
            flags={"read_only": True},
        )
        properties.previous_keys = AAZListType(
            serialized_name="previousKeys",
            flags={"read_only": True},
        )
        properties.provisioning_state = AAZStrType(
            serialized_name="provisioningState",
            flags={"read_only": True},
        )
        properties.rotation_to_latest_key_version_enabled = AAZBoolType(
            serialized_name="rotationToLatestKeyVersionEnabled",
        )

        previous_keys = _schema_disk_encryption_set_read.properties.previous_keys
        previous_keys.Element = AAZObjectType()
        cls._build_schema_key_for_disk_encryption_set_read(previous_keys.Element)

        tags = _schema_disk_encryption_set_read.tags
        tags.Element = AAZStrType()

        _schema.id = cls._schema_disk_encryption_set_read.id
        _schema.identity = cls._schema_disk_encryption_set_read.identity
        _schema.location = cls._schema_disk_encryption_set_read.location
        _schema.name = cls._schema_disk_encryption_set_read.name
        _schema.properties = cls._schema_disk_encryption_set_read.properties
        _schema.tags = cls._schema_disk_encryption_set_read.tags
        _schema.type = cls._schema_disk_encryption_set_read.type

    _schema_key_for_disk_encryption_set_read = None

    @classmethod
    def _build_schema_key_for_disk_encryption_set_read(cls, _schema):
        if cls._schema_key_for_disk_encryption_set_read is not None:
            _schema.key_url = cls._schema_key_for_disk_encryption_set_read.key_url
            _schema.source_vault = cls._schema_key_for_disk_encryption_set_read.source_vault
            return

        cls._schema_key_for_disk_encryption_set_read = _schema_key_for_disk_encryption_set_read = AAZObjectType()

        key_for_disk_encryption_set_read = _schema_key_for_disk_encryption_set_read
        key_for_disk_encryption_set_read.key_url = AAZStrType(
            serialized_name="keyUrl",
            flags={"required": True},
        )
        key_for_disk_encryption_set_read.source_vault = AAZObjectType(
            serialized_name="sourceVault",
        )

        source_vault = _schema_key_for_disk_encryption_set_read.source_vault
        source_vault.id = AAZStrType()

        _schema.key_url = cls._schema_key_for_disk_encryption_set_read.key_url
        _schema.source_vault = cls._schema_key_for_disk_encryption_set_read.source_vault


__all__ = ["Update"]
