# --------------------------------------------------------------------------------------------
# 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(
    "image create",
)
class Create(AAZCommand):
    """Create a custom Virtual Machine Image from managed disks or snapshots.

    :example: Create an image from an existing disk.
        az image create -g MyResourceGroup -n image1 --os-type Linux --source /subscriptions/db5eb68e-73e2-4fa8-b18a-0123456789999/resourceGroups/rg1/providers/Microsoft.Compute/snapshots/s1

    :example: Create an image by capturing an existing generalized virtual machine in the same resource group.
        az image create -g MyResourceGroup -n image1 --source MyVm1
    """

    _aaz_info = {
        "version": "2024-07-01",
        "resources": [
            ["mgmt-plane", "/subscriptions/{}/resourcegroups/{}/providers/microsoft.compute/images/{}", "2024-07-01"],
        ]
    }

    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.image_name = AAZStrArg(
            options=["-n", "--name", "--image-name"],
            help="The name of the image.",
            required=True,
        )
        _args_schema.resource_group = AAZResourceGroupNameArg(
            required=True,
        )

        # define Arg Group "Parameters"

        _args_schema = cls._args_schema
        _args_schema.extended_location = AAZObjectArg(
            options=["--extended-location"],
            arg_group="Parameters",
            help="The extended location of the Image.",
        )
        _args_schema.location = AAZResourceLocationArg(
            arg_group="Parameters",
            help="Resource location",
            required=True,
            fmt=AAZResourceLocationArgFormat(
                resource_group_arg="resource_group",
            ),
        )
        _args_schema.tags = AAZDictArg(
            options=["--tags"],
            arg_group="Parameters",
            help="Resource tags",
        )

        extended_location = cls._args_schema.extended_location
        extended_location.name = AAZStrArg(
            options=["name"],
            help="The name of the extended location.",
        )
        extended_location.type = AAZStrArg(
            options=["type"],
            help="The type of the extended location.",
            enum={"EdgeZone": "EdgeZone"},
        )

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

        # define Arg Group "Properties"

        _args_schema = cls._args_schema
        _args_schema.hyper_v_generation = AAZStrArg(
            options=["--hyper-v-generation"],
            arg_group="Properties",
            help="The hypervisor generation of the Virtual Machine created from the image.",
            default="V1",
            enum={"V1": "V1", "V2": "V2"},
        )
        _args_schema.source_virtual_machine = AAZObjectArg(
            options=["--source-virtual-machine"],
            arg_group="Properties",
            help="The source virtual machine from which Image is created.",
        )
        _args_schema.storage_profile = AAZObjectArg(
            options=["--storage-profile"],
            arg_group="Properties",
            help="Specifies the storage settings for the virtual machine disks.",
        )

        source_virtual_machine = cls._args_schema.source_virtual_machine
        source_virtual_machine.id = AAZStrArg(
            options=["id"],
            help="Resource Id",
        )

        storage_profile = cls._args_schema.storage_profile
        storage_profile.data_disks = AAZListArg(
            options=["data-disks"],
            help="Specifies the parameters that are used to add a data disk to a virtual machine. <br><br> For more information about disks, see [About disks and VHDs for Azure virtual machines](https://learn.microsoft.com/azure/virtual-machines/managed-disks-overview).",
        )
        storage_profile.os_disk = AAZObjectArg(
            options=["os-disk"],
            help="Specifies information about the operating system disk used by the virtual machine. <br><br> For more information about disks, see [About disks and VHDs for Azure virtual machines](https://learn.microsoft.com/azure/virtual-machines/managed-disks-overview).",
        )
        storage_profile.zone_resilient = AAZBoolArg(
            options=["zone-resilient"],
            help="Specifies whether an image is zone resilient or not. Default is false. Zone resilient images can be created only in regions that provide Zone Redundant Storage (ZRS).",
        )

        data_disks = cls._args_schema.storage_profile.data_disks
        data_disks.Element = AAZObjectArg()

        _element = cls._args_schema.storage_profile.data_disks.Element
        _element.blob_uri = AAZStrArg(
            options=["blob-uri"],
            help="The Virtual Hard Disk.",
        )
        _element.caching = AAZStrArg(
            options=["caching"],
            help="Specifies the caching requirements. <br><br> Possible values are: <br><br> **None** <br><br> **ReadOnly** <br><br> **ReadWrite** <br><br> Default: **None for Standard storage. ReadOnly for Premium storage**",
            enum={"None": "None", "ReadOnly": "ReadOnly", "ReadWrite": "ReadWrite"},
        )
        _element.disk_encryption_set = AAZObjectArg(
            options=["disk-encryption-set"],
            help="Specifies the customer managed disk encryption set resource id for the managed image disk.",
        )
        cls._build_args_disk_encryption_set_parameters_create(_element.disk_encryption_set)
        _element.disk_size_gb = AAZIntArg(
            options=["disk-size-gb"],
            help="Specifies the size of empty data disks in gigabytes. This element can be used to overwrite the name of the disk in a virtual machine image. <br><br> This value cannot be larger than 1023 GB",
        )
        _element.lun = AAZIntArg(
            options=["lun"],
            help="Specifies the logical unit number of the data disk. This value is used to identify data disks within the VM and therefore must be unique for each data disk attached to a VM.",
            required=True,
        )
        _element.managed_disk = AAZObjectArg(
            options=["managed-disk"],
            help="The managedDisk.",
        )
        _element.snapshot = AAZObjectArg(
            options=["snapshot"],
            help="The snapshot.",
        )
        _element.storage_account_type = AAZStrArg(
            options=["storage-account-type"],
            help="Specifies the storage account type for the managed disk. NOTE: UltraSSD_LRS can only be used with data disks, it cannot be used with OS Disk.",
            enum={"PremiumV2_LRS": "PremiumV2_LRS", "Premium_LRS": "Premium_LRS", "Premium_ZRS": "Premium_ZRS", "StandardSSD_LRS": "StandardSSD_LRS", "StandardSSD_ZRS": "StandardSSD_ZRS", "Standard_LRS": "Standard_LRS", "UltraSSD_LRS": "UltraSSD_LRS"},
        )

        managed_disk = cls._args_schema.storage_profile.data_disks.Element.managed_disk
        managed_disk.id = AAZStrArg(
            options=["id"],
        )

        snapshot = cls._args_schema.storage_profile.data_disks.Element.snapshot
        snapshot.id = AAZStrArg(
            options=["id"],
        )

        os_disk = cls._args_schema.storage_profile.os_disk
        os_disk.blob_uri = AAZStrArg(
            options=["blob-uri"],
            help="The Virtual Hard Disk.",
        )
        os_disk.caching = AAZStrArg(
            options=["caching"],
            help="Specifies the caching requirements. <br><br> Possible values are: <br><br> **None** <br><br> **ReadOnly** <br><br> **ReadWrite** <br><br> Default: **None for Standard storage. ReadOnly for Premium storage**",
            enum={"None": "None", "ReadOnly": "ReadOnly", "ReadWrite": "ReadWrite"},
        )
        os_disk.disk_encryption_set = AAZObjectArg(
            options=["disk-encryption-set"],
            help="Specifies the customer managed disk encryption set resource id for the managed image disk.",
        )
        cls._build_args_disk_encryption_set_parameters_create(os_disk.disk_encryption_set)
        os_disk.disk_size_gb = AAZIntArg(
            options=["disk-size-gb"],
            help="Specifies the size of empty data disks in gigabytes. This element can be used to overwrite the name of the disk in a virtual machine image. <br><br> This value cannot be larger than 1023 GB",
        )
        os_disk.managed_disk = AAZObjectArg(
            options=["managed-disk"],
            help="The managedDisk.",
        )
        cls._build_args_sub_resource_create(os_disk.managed_disk)
        os_disk.os_state = AAZStrArg(
            options=["os-state"],
            help="The OS State. For managed images, use Generalized.",
            required=True,
            enum={"Generalized": "Generalized", "Specialized": "Specialized"},
        )
        os_disk.os_type = AAZStrArg(
            options=["os-type"],
            help="This property allows you to specify the type of the OS that is included in the disk if creating a VM from a custom image. <br><br> Possible values are: <br><br> **Windows** <br><br> **Linux**",
            required=True,
            enum={"Linux": "Linux", "Windows": "Windows"},
        )
        os_disk.snapshot = AAZObjectArg(
            options=["snapshot"],
            help="The snapshot.",
        )
        os_disk.storage_account_type = AAZStrArg(
            options=["storage-account-type"],
            help="Specifies the storage account type for the managed disk. NOTE: UltraSSD_LRS can only be used with data disks, it cannot be used with OS Disk.",
            enum={"PremiumV2_LRS": "PremiumV2_LRS", "Premium_LRS": "Premium_LRS", "Premium_ZRS": "Premium_ZRS", "StandardSSD_LRS": "StandardSSD_LRS", "StandardSSD_ZRS": "StandardSSD_ZRS", "Standard_LRS": "Standard_LRS", "UltraSSD_LRS": "UltraSSD_LRS"},
        )

        snapshot = cls._args_schema.storage_profile.os_disk.snapshot
        snapshot.id = AAZStrArg(
            options=["id"],
        )
        return cls._args_schema

    _args_disk_encryption_set_parameters_create = None

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

        cls._args_disk_encryption_set_parameters_create = AAZObjectArg()

        disk_encryption_set_parameters_create = cls._args_disk_encryption_set_parameters_create
        disk_encryption_set_parameters_create.id = AAZStrArg(
            options=["id"],
            help="Resource Id",
        )

        _schema.id = cls._args_disk_encryption_set_parameters_create.id

    _args_sub_resource_create = None

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

        cls._args_sub_resource_create = AAZObjectArg()

        sub_resource_create = cls._args_sub_resource_create
        sub_resource_create.id = AAZStrArg(
            options=["id"],
            help="Resource Id",
        )

        _schema.id = cls._args_sub_resource_create.id

    def _execute_operations(self):
        self.pre_operations()
        yield self.ImagesCreateOrUpdate(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 ImagesCreateOrUpdate(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_201,
                    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, 201]:
                return self.client.build_lro_polling(
                    self.ctx.args.no_wait,
                    session,
                    self.on_200_201,
                    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/images/{imageName}",
                **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(
                    "imageName", self.ctx.args.image_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", "2024-07-01",
                    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("extendedLocation", AAZObjectType, ".extended_location")
            _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")

            extended_location = _builder.get(".extendedLocation")
            if extended_location is not None:
                extended_location.set_prop("name", AAZStrType, ".name")
                extended_location.set_prop("type", AAZStrType, ".type")

            properties = _builder.get(".properties")
            if properties is not None:
                properties.set_prop("hyperVGeneration", AAZStrType, ".hyper_v_generation")
                properties.set_prop("sourceVirtualMachine", AAZObjectType, ".source_virtual_machine")
                properties.set_prop("storageProfile", AAZObjectType, ".storage_profile")

            source_virtual_machine = _builder.get(".properties.sourceVirtualMachine")
            if source_virtual_machine is not None:
                source_virtual_machine.set_prop("id", AAZStrType, ".id")

            storage_profile = _builder.get(".properties.storageProfile")
            if storage_profile is not None:
                storage_profile.set_prop("dataDisks", AAZListType, ".data_disks")
                storage_profile.set_prop("osDisk", AAZObjectType, ".os_disk")
                storage_profile.set_prop("zoneResilient", AAZBoolType, ".zone_resilient")

            data_disks = _builder.get(".properties.storageProfile.dataDisks")
            if data_disks is not None:
                data_disks.set_elements(AAZObjectType, ".")

            _elements = _builder.get(".properties.storageProfile.dataDisks[]")
            if _elements is not None:
                _elements.set_prop("blobUri", AAZStrType, ".blob_uri")
                _elements.set_prop("caching", AAZStrType, ".caching")
                _CreateHelper._build_schema_disk_encryption_set_parameters_create(_elements.set_prop("diskEncryptionSet", AAZObjectType, ".disk_encryption_set"))
                _elements.set_prop("diskSizeGB", AAZIntType, ".disk_size_gb")
                _elements.set_prop("lun", AAZIntType, ".lun", typ_kwargs={"flags": {"required": True}})
                _elements.set_prop("managedDisk", AAZObjectType, ".managed_disk")
                _elements.set_prop("snapshot", AAZObjectType, ".snapshot")
                _elements.set_prop("storageAccountType", AAZStrType, ".storage_account_type")

            managed_disk = _builder.get(".properties.storageProfile.dataDisks[].managedDisk")
            if managed_disk is not None:
                managed_disk.set_prop("id", AAZStrType, ".id")

            snapshot = _builder.get(".properties.storageProfile.dataDisks[].snapshot")
            if snapshot is not None:
                snapshot.set_prop("id", AAZStrType, ".id")

            os_disk = _builder.get(".properties.storageProfile.osDisk")
            if os_disk is not None:
                os_disk.set_prop("blobUri", AAZStrType, ".blob_uri")
                os_disk.set_prop("caching", AAZStrType, ".caching")
                _CreateHelper._build_schema_disk_encryption_set_parameters_create(os_disk.set_prop("diskEncryptionSet", AAZObjectType, ".disk_encryption_set"))
                os_disk.set_prop("diskSizeGB", AAZIntType, ".disk_size_gb")
                _CreateHelper._build_schema_sub_resource_create(os_disk.set_prop("managedDisk", AAZObjectType, ".managed_disk"))
                os_disk.set_prop("osState", AAZStrType, ".os_state", typ_kwargs={"flags": {"required": True}})
                os_disk.set_prop("osType", AAZStrType, ".os_type", typ_kwargs={"flags": {"required": True}})
                os_disk.set_prop("snapshot", AAZObjectType, ".snapshot")
                os_disk.set_prop("storageAccountType", AAZStrType, ".storage_account_type")

            snapshot = _builder.get(".properties.storageProfile.osDisk.snapshot")
            if snapshot is not None:
                snapshot.set_prop("id", AAZStrType, ".id")

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

            return self.serialize_content(_content_value)

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

        _schema_on_200_201 = None

        @classmethod
        def _build_schema_on_200_201(cls):
            if cls._schema_on_200_201 is not None:
                return cls._schema_on_200_201

            cls._schema_on_200_201 = AAZObjectType()

            _schema_on_200_201 = cls._schema_on_200_201
            _schema_on_200_201.extended_location = AAZObjectType(
                serialized_name="extendedLocation",
            )
            _schema_on_200_201.id = AAZStrType(
                flags={"read_only": True},
            )
            _schema_on_200_201.location = AAZStrType(
                flags={"required": True},
            )
            _schema_on_200_201.name = AAZStrType(
                flags={"read_only": True},
            )
            _schema_on_200_201.properties = AAZObjectType(
                flags={"client_flatten": True},
            )
            _schema_on_200_201.tags = AAZDictType()
            _schema_on_200_201.type = AAZStrType(
                flags={"read_only": True},
            )

            extended_location = cls._schema_on_200_201.extended_location
            extended_location.name = AAZStrType()
            extended_location.type = AAZStrType()

            properties = cls._schema_on_200_201.properties
            properties.hyper_v_generation = AAZStrType(
                serialized_name="hyperVGeneration",
            )
            properties.provisioning_state = AAZStrType(
                serialized_name="provisioningState",
                flags={"read_only": True},
            )
            properties.source_virtual_machine = AAZObjectType(
                serialized_name="sourceVirtualMachine",
            )
            _CreateHelper._build_schema_sub_resource_read(properties.source_virtual_machine)
            properties.storage_profile = AAZObjectType(
                serialized_name="storageProfile",
            )

            storage_profile = cls._schema_on_200_201.properties.storage_profile
            storage_profile.data_disks = AAZListType(
                serialized_name="dataDisks",
            )
            storage_profile.os_disk = AAZObjectType(
                serialized_name="osDisk",
            )
            storage_profile.zone_resilient = AAZBoolType(
                serialized_name="zoneResilient",
            )

            data_disks = cls._schema_on_200_201.properties.storage_profile.data_disks
            data_disks.Element = AAZObjectType()

            _element = cls._schema_on_200_201.properties.storage_profile.data_disks.Element
            _element.blob_uri = AAZStrType(
                serialized_name="blobUri",
            )
            _element.caching = AAZStrType()
            _element.disk_encryption_set = AAZObjectType(
                serialized_name="diskEncryptionSet",
            )
            _CreateHelper._build_schema_disk_encryption_set_parameters_read(_element.disk_encryption_set)
            _element.disk_size_gb = AAZIntType(
                serialized_name="diskSizeGB",
            )
            _element.lun = AAZIntType(
                flags={"required": True},
            )
            _element.managed_disk = AAZObjectType(
                serialized_name="managedDisk",
            )
            _CreateHelper._build_schema_sub_resource_read(_element.managed_disk)
            _element.snapshot = AAZObjectType()
            _CreateHelper._build_schema_sub_resource_read(_element.snapshot)
            _element.storage_account_type = AAZStrType(
                serialized_name="storageAccountType",
            )

            os_disk = cls._schema_on_200_201.properties.storage_profile.os_disk
            os_disk.blob_uri = AAZStrType(
                serialized_name="blobUri",
            )
            os_disk.caching = AAZStrType()
            os_disk.disk_encryption_set = AAZObjectType(
                serialized_name="diskEncryptionSet",
            )
            _CreateHelper._build_schema_disk_encryption_set_parameters_read(os_disk.disk_encryption_set)
            os_disk.disk_size_gb = AAZIntType(
                serialized_name="diskSizeGB",
            )
            os_disk.managed_disk = AAZObjectType(
                serialized_name="managedDisk",
            )
            _CreateHelper._build_schema_sub_resource_read(os_disk.managed_disk)
            os_disk.os_state = AAZStrType(
                serialized_name="osState",
                flags={"required": True},
            )
            os_disk.os_type = AAZStrType(
                serialized_name="osType",
                flags={"required": True},
            )
            os_disk.snapshot = AAZObjectType()
            _CreateHelper._build_schema_sub_resource_read(os_disk.snapshot)
            os_disk.storage_account_type = AAZStrType(
                serialized_name="storageAccountType",
            )

            tags = cls._schema_on_200_201.tags
            tags.Element = AAZStrType()

            return cls._schema_on_200_201


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

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

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

    _schema_disk_encryption_set_parameters_read = None

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

        cls._schema_disk_encryption_set_parameters_read = _schema_disk_encryption_set_parameters_read = AAZObjectType()

        disk_encryption_set_parameters_read = _schema_disk_encryption_set_parameters_read
        disk_encryption_set_parameters_read.id = AAZStrType()

        _schema.id = cls._schema_disk_encryption_set_parameters_read.id

    _schema_sub_resource_read = None

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

        cls._schema_sub_resource_read = _schema_sub_resource_read = AAZObjectType()

        sub_resource_read = _schema_sub_resource_read
        sub_resource_read.id = AAZStrType()

        _schema.id = cls._schema_sub_resource_read.id


__all__ = ["Create"]
