# --------------------------------------------------------------------------------------------
# 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 create",
)
class Create(AAZCommand):
    """Create a disk encryption set

    :example: Create a disk encryption set
        az disk-encryption-set create --resource-group MyResourceGroup --name MyDiskEncryptionSet --key-url MyKey --source-vault MyVault

    :example: Create a disk encryption set with a system assigned identity.
        az disk-encryption-set create --resource-group MyResourceGroup --name MyDiskEncryptionSet --key-url MyKey --source-vault MyVault --mi-system-assigned

    :example: Create a disk encryption set with a user assigned identity.
        az disk-encryption-set create --resource-group MyResourceGroup --name MyDiskEncryptionSet --key-url MyKey --source-vault MyVault --mi-user-assigned myAssignedId

    :example: Create a disk encryption set with system assigned identity and a user assigned identity.
        az disk-encryption-set create --resource-group MyResourceGroup --name MyDiskEncryptionSet --key-url MyKey --source-vault MyVault --mi-system-assigned --mi-user-assigned myAssignedId

    :example: Create a disk encryption set with multi-tenant application client id to access key vault in a different tenant.
        az disk-encryption-set create --resource-group MyResourceGroup --name MyDiskEncryptionSet --key-url MyKey --source-vault MyVault --federated-client-id myFederatedClientId

    :example: Create a disk encryption set.
        az disk-encryption-set create --resource-group MyResourceGroup --name MyDiskEncryptionSet --key-url MyKey --source-vault MyVault --encryption-type EncryptionAtRestWithPlatformAndCustomerKeys
    """

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

    AZ_SUPPORT_NO_WAIT = 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,
        )
        _args_schema.resource_group = AAZResourceGroupNameArg(
            required=True,
        )
        _args_schema.location = AAZResourceLocationArg(
            help="Resource location",
            required=True,
            fmt=AAZResourceLocationArgFormat(
                resource_group_arg="resource_group",
            ),
        )
        _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.",
        )
        _args_schema.encryption_type = AAZStrArg(
            options=["--encryption-type"],
            help="The type of key used to encrypt the data of the disk. EncryptionAtRestWithCustomerKey: Disk is encrypted at rest with Customer managed key that can be changed and revoked by a customer. EncryptionAtRestWithPlatformAndCustomerKeys: Disk is encrypted at rest with 2 layers of encryption. One of the keys is Customer managed and the other key is Platform managed. ConfidentialVmEncryptedWithCustomerKey: An additional encryption type accepted for confidential VM. Disk is encrypted at rest with Customer managed key.",
            enum={"ConfidentialVmEncryptedWithCustomerKey": "ConfidentialVmEncryptedWithCustomerKey", "EncryptionAtRestWithCustomerKey": "EncryptionAtRestWithCustomerKey", "EncryptionAtRestWithPlatformAndCustomerKeys": "EncryptionAtRestWithPlatformAndCustomerKeys"},
        )
        _args_schema.federated_client_id = AAZStrArg(
            options=["--federated-client-id"],
            help="The federated client id used in cross tenant scenario.",
        )
        _args_schema.enable_auto_key_rotation = AAZBoolArg(
            options=["--auto-rotation", "--enable-auto-key-rotation"],
            help="Enable automatic rotation of keys.",
        )
        _args_schema.tags = AAZDictArg(
            options=["--tags"],
            help="Resource tags",
        )

        tags = cls._args_schema.tags
        tags.Element = AAZStrArg()

        # define Arg Group "Managed Identity"

        _args_schema = cls._args_schema
        _args_schema.identity = AAZObjectArg(
            options=["--identity"],
            arg_group="Managed Identity",
            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.",
        )

        identity = cls._args_schema.identity
        identity.mi_system_assigned = AAZStrArg(
            options=["system-assigned", "mi-system-assigned"],
            help="Set the system managed identity.",
            blank="True",
        )
        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.",
            enum={"None": "None", "SystemAssigned": "SystemAssigned", "SystemAssigned, UserAssigned": "SystemAssigned, UserAssigned", "UserAssigned": "UserAssigned"},
        )
        identity.mi_user_assigned = AAZListArg(
            options=["user-assigned", "mi-user-assigned"],
            help="Set the user managed identities.",
            blank=[],
        )
        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}'.",
        )

        mi_user_assigned = cls._args_schema.identity.mi_user_assigned
        mi_user_assigned.Element = AAZStrArg()

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

    def _execute_operations(self):
        self.pre_operations()
        yield self.DiskEncryptionSetsCreateOrUpdate(ctx=self.ctx)()
        self.post_operations()

    @register_callback
    def pre_operations(self):
        pass

    @register_callback
    def post_operations(self):
        pass

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

    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,
                typ=AAZObjectType,
                typ_kwargs={"flags": {"required": True, "client_flatten": True}}
            )
            _builder.set_prop("identity", AAZIdentityObjectType, ".identity")
            _builder.set_prop("location", AAZStrType, ".location", typ_kwargs={"flags": {"required": True}})
            _builder.set_prop("properties", AAZObjectType, typ_kwargs={"flags": {"client_flatten": True}})
            _builder.set_prop("tags", AAZDictType, ".tags")

            identity = _builder.get(".identity")
            if identity is not None:
                identity.set_prop("type", AAZStrType, ".type")
                identity.set_prop("userAssignedIdentities", AAZDictType, ".user_assigned_identities")
                identity.set_prop("userAssigned", AAZListType, ".mi_user_assigned", typ_kwargs={"flags": {"action": "create"}})
                identity.set_prop("systemAssigned", AAZStrType, ".mi_system_assigned", typ_kwargs={"flags": {"action": "create"}})

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

            user_assigned = _builder.get(".identity.userAssigned")
            if user_assigned is not None:
                user_assigned.set_elements(AAZStrType, ".")

            properties = _builder.get(".properties")
            if properties is not None:
                properties.set_prop("activeKey", AAZObjectType)
                properties.set_prop("encryptionType", AAZStrType, ".encryption_type")
                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")

            tags = _builder.get(".tags")
            if tags is not None:
                tags.set_elements(AAZStrType, ".")

            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()
            _CreateHelper._build_schema_disk_encryption_set_read(cls._schema_on_200)

            return cls._schema_on_200


class _CreateHelper:
    """Helper class for Create"""

    _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__ = ["Create"]
