# --------------------------------------------------------------------------------------------
# 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 *


class UpdatePatch(AAZCommand):
    """Update disk config.

    :example: Update disk size.
        az disk config update --name MyManagedDisk --resource-group MyResourceGroup --size-gb 20
    """

    _aaz_info = {
        "version": "2023-04-02",
        "resources": [
            ["mgmt-plane", "/subscriptions/{}/resourcegroups/{}/providers/microsoft.compute/disks/{}", "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_name = AAZStrArg(
            options=["-n", "--name", "--disk-name"],
            help="The name of the managed disk that is being created. The name can't be changed after the disk is created. Supported characters for the name are a-z, A-Z, 0-9, _ and -. The maximum name length is 80 characters.",
            required=True,
            id_part="name",
        )
        _args_schema.resource_group = AAZResourceGroupNameArg(
            required=True,
        )

        # define Arg Group "Disk"

        # define Arg Group "Properties"

        _args_schema = cls._args_schema
        _args_schema.disk_size_gb = AAZIntArg(
            options=["-z", "--size-gb", "--disk-size-gb"],
            arg_group="Properties",
            help="If creationData.createOption is Empty, this field is mandatory and it indicates the size of the disk to create. If this field is present for updates or creation with other options, it indicates a resize. Resizes are only allowed if the disk is not attached to a running VM, and can only increase the disk's size.",
        )
        return cls._args_schema

    _args_source_vault_update = None

    @classmethod
    def _build_args_source_vault_update(cls, _schema):
        if cls._args_source_vault_update is not None:
            _schema.id = cls._args_source_vault_update.id
            return

        cls._args_source_vault_update = AAZObjectArg()

        source_vault_update = cls._args_source_vault_update
        source_vault_update.id = AAZStrArg(
            options=["id"],
            help="Resource Id",
        )

        _schema.id = cls._args_source_vault_update.id

    def _execute_operations(self):
        self.pre_operations()
        yield self.DisksUpdate(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 DisksUpdate(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/disks/{diskName}",
                **self.url_parameters
            )

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

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

        @property
        def url_parameters(self):
            parameters = {
                **self.serialize_url_param(
                    "diskName", self.ctx.args.disk_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("properties", AAZObjectType, typ_kwargs={"flags": {"client_flatten": True}})

            properties = _builder.get(".properties")
            if properties is not None:
                properties.set_prop("diskSizeGB", AAZIntType, ".disk_size_gb")

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

            return cls._schema_on_200


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

    @classmethod
    def _build_schema_source_vault_update(cls, _builder):
        if _builder is None:
            return
        _builder.set_prop("id", AAZStrType, ".id")

    _schema_disk_read = None

    @classmethod
    def _build_schema_disk_read(cls, _schema):
        if cls._schema_disk_read is not None:
            _schema.extended_location = cls._schema_disk_read.extended_location
            _schema.id = cls._schema_disk_read.id
            _schema.location = cls._schema_disk_read.location
            _schema.managed_by = cls._schema_disk_read.managed_by
            _schema.managed_by_extended = cls._schema_disk_read.managed_by_extended
            _schema.name = cls._schema_disk_read.name
            _schema.properties = cls._schema_disk_read.properties
            _schema.sku = cls._schema_disk_read.sku
            _schema.tags = cls._schema_disk_read.tags
            _schema.type = cls._schema_disk_read.type
            _schema.zones = cls._schema_disk_read.zones
            return

        cls._schema_disk_read = _schema_disk_read = AAZObjectType()

        disk_read = _schema_disk_read
        disk_read.extended_location = AAZObjectType(
            serialized_name="extendedLocation",
        )
        disk_read.id = AAZStrType(
            flags={"read_only": True},
        )
        disk_read.location = AAZStrType(
            flags={"required": True},
        )
        disk_read.managed_by = AAZStrType(
            serialized_name="managedBy",
            flags={"read_only": True},
        )
        disk_read.managed_by_extended = AAZListType(
            serialized_name="managedByExtended",
            flags={"read_only": True},
        )
        disk_read.name = AAZStrType(
            flags={"read_only": True},
        )
        disk_read.properties = AAZObjectType(
            flags={"client_flatten": True},
        )
        disk_read.sku = AAZObjectType()
        disk_read.tags = AAZDictType()
        disk_read.type = AAZStrType(
            flags={"read_only": True},
        )
        disk_read.zones = AAZListType()

        extended_location = _schema_disk_read.extended_location
        extended_location.name = AAZStrType()
        extended_location.type = AAZStrType()

        managed_by_extended = _schema_disk_read.managed_by_extended
        managed_by_extended.Element = AAZStrType()

        properties = _schema_disk_read.properties
        properties.last_ownership_update_time = AAZStrType(
            serialized_name="LastOwnershipUpdateTime",
            flags={"read_only": True},
        )
        properties.bursting_enabled = AAZBoolType(
            serialized_name="burstingEnabled",
        )
        properties.bursting_enabled_time = AAZStrType(
            serialized_name="burstingEnabledTime",
            flags={"read_only": True},
        )
        properties.completion_percent = AAZFloatType(
            serialized_name="completionPercent",
        )
        properties.creation_data = AAZObjectType(
            serialized_name="creationData",
            flags={"required": True},
        )
        properties.data_access_auth_mode = AAZStrType(
            serialized_name="dataAccessAuthMode",
        )
        properties.disk_access_id = AAZStrType(
            serialized_name="diskAccessId",
        )
        properties.disk_iops_read_only = AAZIntType(
            serialized_name="diskIOPSReadOnly",
        )
        properties.disk_iops_read_write = AAZIntType(
            serialized_name="diskIOPSReadWrite",
        )
        properties.disk_m_bps_read_only = AAZIntType(
            serialized_name="diskMBpsReadOnly",
        )
        properties.disk_m_bps_read_write = AAZIntType(
            serialized_name="diskMBpsReadWrite",
        )
        properties.disk_size_bytes = AAZIntType(
            serialized_name="diskSizeBytes",
            flags={"read_only": True},
        )
        properties.disk_size_gb = AAZIntType(
            serialized_name="diskSizeGB",
        )
        properties.disk_state = AAZStrType(
            serialized_name="diskState",
            flags={"read_only": True},
        )
        properties.encryption = AAZObjectType()
        properties.encryption_settings_collection = AAZObjectType(
            serialized_name="encryptionSettingsCollection",
        )
        properties.hyper_v_generation = AAZStrType(
            serialized_name="hyperVGeneration",
        )
        properties.max_shares = AAZIntType(
            serialized_name="maxShares",
        )
        properties.network_access_policy = AAZStrType(
            serialized_name="networkAccessPolicy",
        )
        properties.optimized_for_frequent_attach = AAZBoolType(
            serialized_name="optimizedForFrequentAttach",
        )
        properties.os_type = AAZStrType(
            serialized_name="osType",
        )
        properties.property_updates_in_progress = AAZObjectType(
            serialized_name="propertyUpdatesInProgress",
            flags={"read_only": True},
        )
        properties.provisioning_state = AAZStrType(
            serialized_name="provisioningState",
            flags={"read_only": True},
        )
        properties.public_network_access = AAZStrType(
            serialized_name="publicNetworkAccess",
        )
        properties.purchase_plan = AAZObjectType(
            serialized_name="purchasePlan",
        )
        properties.security_profile = AAZObjectType(
            serialized_name="securityProfile",
        )
        properties.share_info = AAZListType(
            serialized_name="shareInfo",
            flags={"read_only": True},
        )
        properties.supported_capabilities = AAZObjectType(
            serialized_name="supportedCapabilities",
        )
        properties.supports_hibernation = AAZBoolType(
            serialized_name="supportsHibernation",
        )
        properties.tier = AAZStrType()
        properties.time_created = AAZStrType(
            serialized_name="timeCreated",
            flags={"read_only": True},
        )
        properties.unique_id = AAZStrType(
            serialized_name="uniqueId",
            flags={"read_only": True},
        )

        creation_data = _schema_disk_read.properties.creation_data
        creation_data.create_option = AAZStrType(
            serialized_name="createOption",
            flags={"required": True},
        )
        creation_data.elastic_san_resource_id = AAZStrType(
            serialized_name="elasticSanResourceId",
        )
        creation_data.gallery_image_reference = AAZObjectType(
            serialized_name="galleryImageReference",
        )
        cls._build_schema_image_disk_reference_read(creation_data.gallery_image_reference)
        creation_data.image_reference = AAZObjectType(
            serialized_name="imageReference",
        )
        cls._build_schema_image_disk_reference_read(creation_data.image_reference)
        creation_data.logical_sector_size = AAZIntType(
            serialized_name="logicalSectorSize",
        )
        creation_data.performance_plus = AAZBoolType(
            serialized_name="performancePlus",
        )
        creation_data.security_data_uri = AAZStrType(
            serialized_name="securityDataUri",
        )
        creation_data.source_resource_id = AAZStrType(
            serialized_name="sourceResourceId",
        )
        creation_data.source_unique_id = AAZStrType(
            serialized_name="sourceUniqueId",
            flags={"read_only": True},
        )
        creation_data.source_uri = AAZStrType(
            serialized_name="sourceUri",
        )
        creation_data.storage_account_id = AAZStrType(
            serialized_name="storageAccountId",
        )
        creation_data.upload_size_bytes = AAZIntType(
            serialized_name="uploadSizeBytes",
        )

        encryption = _schema_disk_read.properties.encryption
        encryption.disk_encryption_set_id = AAZStrType(
            serialized_name="diskEncryptionSetId",
        )
        encryption.type = AAZStrType()

        encryption_settings_collection = _schema_disk_read.properties.encryption_settings_collection
        encryption_settings_collection.enabled = AAZBoolType(
            flags={"required": True},
        )
        encryption_settings_collection.encryption_settings = AAZListType(
            serialized_name="encryptionSettings",
        )
        encryption_settings_collection.encryption_settings_version = AAZStrType(
            serialized_name="encryptionSettingsVersion",
        )

        encryption_settings = _schema_disk_read.properties.encryption_settings_collection.encryption_settings
        encryption_settings.Element = AAZObjectType()

        _element = _schema_disk_read.properties.encryption_settings_collection.encryption_settings.Element
        _element.disk_encryption_key = AAZObjectType(
            serialized_name="diskEncryptionKey",
        )
        _element.key_encryption_key = AAZObjectType(
            serialized_name="keyEncryptionKey",
        )

        disk_encryption_key = _schema_disk_read.properties.encryption_settings_collection.encryption_settings.Element.disk_encryption_key
        disk_encryption_key.secret_url = AAZStrType(
            serialized_name="secretUrl",
            flags={"required": True},
        )
        disk_encryption_key.source_vault = AAZObjectType(
            serialized_name="sourceVault",
            flags={"required": True},
        )
        cls._build_schema_source_vault_read(disk_encryption_key.source_vault)

        key_encryption_key = _schema_disk_read.properties.encryption_settings_collection.encryption_settings.Element.key_encryption_key
        key_encryption_key.key_url = AAZStrType(
            serialized_name="keyUrl",
            flags={"required": True},
        )
        key_encryption_key.source_vault = AAZObjectType(
            serialized_name="sourceVault",
            flags={"required": True},
        )
        cls._build_schema_source_vault_read(key_encryption_key.source_vault)

        property_updates_in_progress = _schema_disk_read.properties.property_updates_in_progress
        property_updates_in_progress.target_tier = AAZStrType(
            serialized_name="targetTier",
        )

        purchase_plan = _schema_disk_read.properties.purchase_plan
        purchase_plan.name = AAZStrType(
            flags={"required": True},
        )
        purchase_plan.product = AAZStrType(
            flags={"required": True},
        )
        purchase_plan.promotion_code = AAZStrType(
            serialized_name="promotionCode",
        )
        purchase_plan.publisher = AAZStrType(
            flags={"required": True},
        )

        security_profile = _schema_disk_read.properties.security_profile
        security_profile.secure_vm_disk_encryption_set_id = AAZStrType(
            serialized_name="secureVMDiskEncryptionSetId",
        )
        security_profile.security_type = AAZStrType(
            serialized_name="securityType",
        )

        share_info = _schema_disk_read.properties.share_info
        share_info.Element = AAZObjectType()

        _element = _schema_disk_read.properties.share_info.Element
        _element.vm_uri = AAZStrType(
            serialized_name="vmUri",
            flags={"read_only": True},
        )

        supported_capabilities = _schema_disk_read.properties.supported_capabilities
        supported_capabilities.accelerated_network = AAZBoolType(
            serialized_name="acceleratedNetwork",
        )
        supported_capabilities.architecture = AAZStrType()
        supported_capabilities.disk_controller_types = AAZStrType(
            serialized_name="diskControllerTypes",
        )

        sku = _schema_disk_read.sku
        sku.name = AAZStrType()
        sku.tier = AAZStrType(
            flags={"read_only": True},
        )

        tags = _schema_disk_read.tags
        tags.Element = AAZStrType()

        zones = _schema_disk_read.zones
        zones.Element = AAZStrType()

        _schema.extended_location = cls._schema_disk_read.extended_location
        _schema.id = cls._schema_disk_read.id
        _schema.location = cls._schema_disk_read.location
        _schema.managed_by = cls._schema_disk_read.managed_by
        _schema.managed_by_extended = cls._schema_disk_read.managed_by_extended
        _schema.name = cls._schema_disk_read.name
        _schema.properties = cls._schema_disk_read.properties
        _schema.sku = cls._schema_disk_read.sku
        _schema.tags = cls._schema_disk_read.tags
        _schema.type = cls._schema_disk_read.type
        _schema.zones = cls._schema_disk_read.zones

    _schema_image_disk_reference_read = None

    @classmethod
    def _build_schema_image_disk_reference_read(cls, _schema):
        if cls._schema_image_disk_reference_read is not None:
            _schema.community_gallery_image_id = cls._schema_image_disk_reference_read.community_gallery_image_id
            _schema.id = cls._schema_image_disk_reference_read.id
            _schema.lun = cls._schema_image_disk_reference_read.lun
            _schema.shared_gallery_image_id = cls._schema_image_disk_reference_read.shared_gallery_image_id
            return

        cls._schema_image_disk_reference_read = _schema_image_disk_reference_read = AAZObjectType()

        image_disk_reference_read = _schema_image_disk_reference_read
        image_disk_reference_read.community_gallery_image_id = AAZStrType(
            serialized_name="communityGalleryImageId",
        )
        image_disk_reference_read.id = AAZStrType()
        image_disk_reference_read.lun = AAZIntType()
        image_disk_reference_read.shared_gallery_image_id = AAZStrType(
            serialized_name="sharedGalleryImageId",
        )

        _schema.community_gallery_image_id = cls._schema_image_disk_reference_read.community_gallery_image_id
        _schema.id = cls._schema_image_disk_reference_read.id
        _schema.lun = cls._schema_image_disk_reference_read.lun
        _schema.shared_gallery_image_id = cls._schema_image_disk_reference_read.shared_gallery_image_id

    _schema_source_vault_read = None

    @classmethod
    def _build_schema_source_vault_read(cls, _schema):
        if cls._schema_source_vault_read is not None:
            _schema.id = cls._schema_source_vault_read.id
            return

        cls._schema_source_vault_read = _schema_source_vault_read = AAZObjectType()

        source_vault_read = _schema_source_vault_read
        source_vault_read.id = AAZStrType()

        _schema.id = cls._schema_source_vault_read.id


__all__ = ["UpdatePatch"]
