# --------------------------------------------------------------------------------------------
# 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(
    "vm availability-set update",
)
class Update(AAZCommand):
    """Update an Azure Availability Set.

    :example: Update an availability set
        az vm availability-set update -n MyAvSet -g MyResourceGroup

    :example: Update an availability set tag.
        az vm availability-set update -n MyAvSet -g MyResourceGroup --set tags.foo=value

    :example: Remove an availability set tag.
        az vm availability-set update -n MyAvSet -g MyResourceGroup --remove tags.foo
    """

    _aaz_info = {
        "version": "2025-04-01",
        "resources": [
            ["mgmt-plane", "/subscriptions/{}/resourcegroups/{}/providers/microsoft.compute/availabilitysets/{}", "2025-04-01"],
        ]
    }

    AZ_SUPPORT_GENERIC_UPDATE = True

    def _handler(self, command_args):
        super()._handler(command_args)
        self._execute_operations()
        return 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.availability_set_name = AAZStrArg(
            options=["-n", "--name", "--availability-set-name"],
            help="The name of the availability set.",
            required=True,
            id_part="name",
        )
        _args_schema.resource_group = AAZResourceGroupNameArg(
            required=True,
        )
        _args_schema.platform_fault_domain_count = AAZIntArg(
            options=["--platform-fault-domain-count"],
            help="Fault Domain count.",
            nullable=True,
        )
        _args_schema.proximity_placement_group = AAZObjectArg(
            options=["--proximity-placement-group"],
            help="Specifies information about the proximity placement group that the availability set should be assigned to. <br><br>Minimum api-version: 2018-04-01.",
            nullable=True,
        )
        cls._build_args_sub_resource_update(_args_schema.proximity_placement_group)
        _args_schema.additional_scheduled_events = AAZBoolArg(
            options=["--additional-events", "--additional-scheduled-events"],
            help="The configuration parameter used while creating event grid and resource graph scheduled event setting.",
            nullable=True,
        )
        _args_schema.enable_user_reboot_scheduled_events = AAZBoolArg(
            options=["--enable-reboot", "--enable-user-reboot-scheduled-events"],
            help="The configuration parameter used while publishing scheduled events additional publishing targets.",
            nullable=True,
        )
        _args_schema.enable_user_redeploy_scheduled_events = AAZBoolArg(
            options=["--enable-redeploy", "--enable-user-redeploy-scheduled-events"],
            help="The configuration parameter used while creating user initiated redeploy scheduled event setting creation.",
            nullable=True,
        )
        _args_schema.sku = AAZObjectArg(
            options=["--sku"],
            help="Sku of the availability set, only name is required to be set. See AvailabilitySetSkuTypes for possible set of values. Use 'Aligned' for virtual machines with managed disks and 'Classic' for virtual machines with unmanaged disks. Default value is 'Classic'.",
            nullable=True,
        )

        sku = cls._args_schema.sku
        sku.capacity = AAZIntArg(
            options=["capacity"],
            help="Specifies the number of virtual machines in the scale set.",
            nullable=True,
        )
        sku.name = AAZStrArg(
            options=["name"],
            help="The sku name.",
            nullable=True,
        )
        sku.tier = AAZStrArg(
            options=["tier"],
            help="Specifies the tier of virtual machines in a scale set.<br /><br /> Possible Values:<br /><br /> **Standard**<br /><br /> **Basic**",
            nullable=True,
        )

        # define Arg Group "AllInstancesDown"

        _args_schema = cls._args_schema
        _args_schema.enable_all_instance_down = AAZBoolArg(
            options=["--all-instance-down", "--enable-all-instance-down"],
            arg_group="AllInstancesDown",
            help="Specify if Scheduled Events should be auto-approved when all instances are down. Its default value is true",
            nullable=True,
        )

        # define Arg Group "EventGridAndResourceGraph"

        _args_schema = cls._args_schema
        _args_schema.scheduled_events_api_version = AAZStrArg(
            options=["--se-api-version", "--scheduled-events-api-version"],
            arg_group="EventGridAndResourceGraph",
            help="Specify the api-version to determine which Scheduled Events configuration schema version will be delivered.",
            nullable=True,
        )

        # define Arg Group "Properties"
        return cls._args_schema

    _args_sub_resource_update = None

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

        cls._args_sub_resource_update = AAZObjectArg(
            nullable=True,
        )

        sub_resource_update = cls._args_sub_resource_update
        sub_resource_update.id = AAZStrArg(
            options=["id"],
            help="Resource Id",
            nullable=True,
        )

        _schema.id = cls._args_sub_resource_update.id

    def _execute_operations(self):
        self.pre_operations()
        self.AvailabilitySetsGet(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)
        self.AvailabilitySetsCreateOrUpdate(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 AvailabilitySetsGet(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/availabilitySets/{availabilitySetName}",
                **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(
                    "availabilitySetName", self.ctx.args.availability_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", "2025-04-01",
                    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_availability_set_read(cls._schema_on_200)

            return cls._schema_on_200

    class AvailabilitySetsCreateOrUpdate(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/availabilitySets/{availabilitySetName}",
                **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(
                    "availabilitySetName", self.ctx.args.availability_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", "2025-04-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,
                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_availability_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("properties", AAZObjectType, typ_kwargs={"flags": {"client_flatten": True}})
            _builder.set_prop("sku", AAZObjectType, ".sku")

            properties = _builder.get(".properties")
            if properties is not None:
                properties.set_prop("platformFaultDomainCount", AAZIntType, ".platform_fault_domain_count")
                _UpdateHelper._build_schema_sub_resource_update(properties.set_prop("proximityPlacementGroup", AAZObjectType, ".proximity_placement_group"))
                properties.set_prop("scheduledEventsPolicy", AAZObjectType)

            scheduled_events_policy = _builder.get(".properties.scheduledEventsPolicy")
            if scheduled_events_policy is not None:
                scheduled_events_policy.set_prop("allInstancesDown", AAZObjectType)
                scheduled_events_policy.set_prop("scheduledEventsAdditionalPublishingTargets", AAZObjectType)
                scheduled_events_policy.set_prop("userInitiatedReboot", AAZObjectType)
                scheduled_events_policy.set_prop("userInitiatedRedeploy", AAZObjectType)

            all_instances_down = _builder.get(".properties.scheduledEventsPolicy.allInstancesDown")
            if all_instances_down is not None:
                all_instances_down.set_prop("automaticallyApprove", AAZBoolType, ".enable_all_instance_down")

            scheduled_events_additional_publishing_targets = _builder.get(".properties.scheduledEventsPolicy.scheduledEventsAdditionalPublishingTargets")
            if scheduled_events_additional_publishing_targets is not None:
                scheduled_events_additional_publishing_targets.set_prop("eventGridAndResourceGraph", AAZObjectType)

            event_grid_and_resource_graph = _builder.get(".properties.scheduledEventsPolicy.scheduledEventsAdditionalPublishingTargets.eventGridAndResourceGraph")
            if event_grid_and_resource_graph is not None:
                event_grid_and_resource_graph.set_prop("enable", AAZBoolType, ".additional_scheduled_events")
                event_grid_and_resource_graph.set_prop("scheduledEventsApiVersion", AAZStrType, ".scheduled_events_api_version")

            user_initiated_reboot = _builder.get(".properties.scheduledEventsPolicy.userInitiatedReboot")
            if user_initiated_reboot is not None:
                user_initiated_reboot.set_prop("automaticallyApprove", AAZBoolType, ".enable_user_reboot_scheduled_events")

            user_initiated_redeploy = _builder.get(".properties.scheduledEventsPolicy.userInitiatedRedeploy")
            if user_initiated_redeploy is not None:
                user_initiated_redeploy.set_prop("automaticallyApprove", AAZBoolType, ".enable_user_redeploy_scheduled_events")

            sku = _builder.get(".sku")
            if sku is not None:
                sku.set_prop("capacity", AAZIntType, ".capacity")
                sku.set_prop("name", AAZStrType, ".name")
                sku.set_prop("tier", AAZStrType, ".tier")

            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"""

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

    _schema_availability_set_read = None

    @classmethod
    def _build_schema_availability_set_read(cls, _schema):
        if cls._schema_availability_set_read is not None:
            _schema.id = cls._schema_availability_set_read.id
            _schema.location = cls._schema_availability_set_read.location
            _schema.name = cls._schema_availability_set_read.name
            _schema.properties = cls._schema_availability_set_read.properties
            _schema.sku = cls._schema_availability_set_read.sku
            _schema.system_data = cls._schema_availability_set_read.system_data
            _schema.tags = cls._schema_availability_set_read.tags
            _schema.type = cls._schema_availability_set_read.type
            return

        cls._schema_availability_set_read = _schema_availability_set_read = AAZObjectType()

        availability_set_read = _schema_availability_set_read
        availability_set_read.id = AAZStrType(
            flags={"read_only": True},
        )
        availability_set_read.location = AAZStrType(
            flags={"required": True},
        )
        availability_set_read.name = AAZStrType(
            flags={"read_only": True},
        )
        availability_set_read.properties = AAZObjectType(
            flags={"client_flatten": True},
        )
        availability_set_read.sku = AAZObjectType()
        availability_set_read.system_data = AAZObjectType(
            serialized_name="systemData",
            flags={"read_only": True},
        )
        availability_set_read.tags = AAZDictType()
        availability_set_read.type = AAZStrType(
            flags={"read_only": True},
        )

        properties = _schema_availability_set_read.properties
        properties.platform_fault_domain_count = AAZIntType(
            serialized_name="platformFaultDomainCount",
        )
        properties.platform_update_domain_count = AAZIntType(
            serialized_name="platformUpdateDomainCount",
        )
        properties.proximity_placement_group = AAZObjectType(
            serialized_name="proximityPlacementGroup",
        )
        cls._build_schema_sub_resource_read(properties.proximity_placement_group)
        properties.scheduled_events_policy = AAZObjectType(
            serialized_name="scheduledEventsPolicy",
        )
        properties.statuses = AAZListType(
            flags={"read_only": True},
        )
        properties.virtual_machine_scale_set_migration_info = AAZObjectType(
            serialized_name="virtualMachineScaleSetMigrationInfo",
            flags={"read_only": True},
        )
        properties.virtual_machines = AAZListType(
            serialized_name="virtualMachines",
        )

        scheduled_events_policy = _schema_availability_set_read.properties.scheduled_events_policy
        scheduled_events_policy.all_instances_down = AAZObjectType(
            serialized_name="allInstancesDown",
        )
        scheduled_events_policy.scheduled_events_additional_publishing_targets = AAZObjectType(
            serialized_name="scheduledEventsAdditionalPublishingTargets",
        )
        scheduled_events_policy.user_initiated_reboot = AAZObjectType(
            serialized_name="userInitiatedReboot",
        )
        scheduled_events_policy.user_initiated_redeploy = AAZObjectType(
            serialized_name="userInitiatedRedeploy",
        )

        all_instances_down = _schema_availability_set_read.properties.scheduled_events_policy.all_instances_down
        all_instances_down.automatically_approve = AAZBoolType(
            serialized_name="automaticallyApprove",
        )

        scheduled_events_additional_publishing_targets = _schema_availability_set_read.properties.scheduled_events_policy.scheduled_events_additional_publishing_targets
        scheduled_events_additional_publishing_targets.event_grid_and_resource_graph = AAZObjectType(
            serialized_name="eventGridAndResourceGraph",
        )

        event_grid_and_resource_graph = _schema_availability_set_read.properties.scheduled_events_policy.scheduled_events_additional_publishing_targets.event_grid_and_resource_graph
        event_grid_and_resource_graph.enable = AAZBoolType()
        event_grid_and_resource_graph.scheduled_events_api_version = AAZStrType(
            serialized_name="scheduledEventsApiVersion",
        )

        user_initiated_reboot = _schema_availability_set_read.properties.scheduled_events_policy.user_initiated_reboot
        user_initiated_reboot.automatically_approve = AAZBoolType(
            serialized_name="automaticallyApprove",
        )

        user_initiated_redeploy = _schema_availability_set_read.properties.scheduled_events_policy.user_initiated_redeploy
        user_initiated_redeploy.automatically_approve = AAZBoolType(
            serialized_name="automaticallyApprove",
        )

        statuses = _schema_availability_set_read.properties.statuses
        statuses.Element = AAZObjectType()

        _element = _schema_availability_set_read.properties.statuses.Element
        _element.code = AAZStrType()
        _element.display_status = AAZStrType(
            serialized_name="displayStatus",
        )
        _element.level = AAZStrType()
        _element.message = AAZStrType()
        _element.time = AAZStrType()

        virtual_machine_scale_set_migration_info = _schema_availability_set_read.properties.virtual_machine_scale_set_migration_info
        virtual_machine_scale_set_migration_info.default_virtual_machine_scale_set_info = AAZObjectType(
            serialized_name="defaultVirtualMachineScaleSetInfo",
            flags={"read_only": True},
        )
        virtual_machine_scale_set_migration_info.migrate_to_virtual_machine_scale_set = AAZObjectType(
            serialized_name="migrateToVirtualMachineScaleSet",
            flags={"read_only": True},
        )
        cls._build_schema_sub_resource_read(virtual_machine_scale_set_migration_info.migrate_to_virtual_machine_scale_set)

        default_virtual_machine_scale_set_info = _schema_availability_set_read.properties.virtual_machine_scale_set_migration_info.default_virtual_machine_scale_set_info
        default_virtual_machine_scale_set_info.constrained_maximum_capacity = AAZBoolType(
            serialized_name="constrainedMaximumCapacity",
            flags={"read_only": True},
        )
        default_virtual_machine_scale_set_info.default_virtual_machine_scale_set = AAZObjectType(
            serialized_name="defaultVirtualMachineScaleSet",
            flags={"read_only": True},
        )
        cls._build_schema_sub_resource_read(default_virtual_machine_scale_set_info.default_virtual_machine_scale_set)

        virtual_machines = _schema_availability_set_read.properties.virtual_machines
        virtual_machines.Element = AAZObjectType()
        cls._build_schema_sub_resource_read(virtual_machines.Element)

        sku = _schema_availability_set_read.sku
        sku.capacity = AAZIntType()
        sku.name = AAZStrType()
        sku.tier = AAZStrType()

        system_data = _schema_availability_set_read.system_data
        system_data.created_at = AAZStrType(
            serialized_name="createdAt",
        )
        system_data.created_by = AAZStrType(
            serialized_name="createdBy",
        )
        system_data.created_by_type = AAZStrType(
            serialized_name="createdByType",
        )
        system_data.last_modified_at = AAZStrType(
            serialized_name="lastModifiedAt",
        )
        system_data.last_modified_by = AAZStrType(
            serialized_name="lastModifiedBy",
        )
        system_data.last_modified_by_type = AAZStrType(
            serialized_name="lastModifiedByType",
        )

        tags = _schema_availability_set_read.tags
        tags.Element = AAZStrType()

        _schema.id = cls._schema_availability_set_read.id
        _schema.location = cls._schema_availability_set_read.location
        _schema.name = cls._schema_availability_set_read.name
        _schema.properties = cls._schema_availability_set_read.properties
        _schema.sku = cls._schema_availability_set_read.sku
        _schema.system_data = cls._schema_availability_set_read.system_data
        _schema.tags = cls._schema_availability_set_read.tags
        _schema.type = cls._schema_availability_set_read.type

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