# --------------------------------------------------------------------------------------------
# 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(
    "search service update",
)
class Update(AAZCommand):
    """Update a search service in the given resource group. If the search service already exists, all properties will be updated with the given values.
    """

    _aaz_info = {
        "version": "2022-09-01",
        "resources": [
            ["mgmt-plane", "/subscriptions/{}/resourcegroups/{}/providers/microsoft.search/searchservices/{}", "2022-09-01"],
        ]
    }

    AZ_SUPPORT_NO_WAIT = True

    AZ_SUPPORT_GENERIC_UPDATE = True

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

    _args_schema = None

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

        # define Arg Group ""

        _args_schema = cls._args_schema
        _args_schema.resource_group = AAZResourceGroupNameArg(
            required=True,
        )
        _args_schema.search_service_name = AAZStrArg(
            options=["-n", "--name", "--search-service-name"],
            help="The name of the Azure Cognitive Search service associated with the specified resource group.",
            required=True,
            id_part="name",
        )

        # define Arg Group "AuthOptions"

        _args_schema = cls._args_schema
        _args_schema.aad_auth_failure_mode = AAZStrArg(
            options=["--aad-auth-failure-mode"],
            arg_group="AuthOptions",
            help="Describes what response the data plane API of a Search service would send for requests that failed authentication.",
            nullable=True,
            enum={"http401WithBearerChallenge": "http401WithBearerChallenge", "http403": "http403"},
        )
        _args_schema.api_key_only = AAZObjectArg(
            options=["--api-key-only"],
            arg_group="AuthOptions",
            help="Indicates that only the API key needs to be used for authentication.",
            nullable=True,
            blank={},
        )

        # define Arg Group "Identity"

        _args_schema = cls._args_schema
        _args_schema.identity_type = AAZStrArg(
            options=["--identity-type"],
            arg_group="Identity",
            help="The identity type.",
            enum={"None": "None", "SystemAssigned": "SystemAssigned"},
        )

        # define Arg Group "NetworkRuleSet"

        _args_schema = cls._args_schema
        _args_schema.ip_rules_internal = AAZListArg(
            options=["--ip-rules-internal"],
            arg_group="NetworkRuleSet",
            help="A list of IP restriction rules that defines the inbound network(s) with allowing access to the search service endpoint. At the meantime, all other public IP networks are blocked by the firewall. These restriction rules are applied only when the 'publicNetworkAccess' of the search service is 'enabled'; otherwise, traffic over public interface is not allowed even with any public IP rules, and private endpoint connections would be the exclusive access method.",
            nullable=True,
        )

        ip_rules_internal = cls._args_schema.ip_rules_internal
        ip_rules_internal.Element = AAZObjectArg(
            nullable=True,
        )

        _element = cls._args_schema.ip_rules_internal.Element
        _element.value = AAZStrArg(
            options=["value"],
            help="Value corresponding to a single IPv4 address (eg., 123.1.2.3) or an IP range in CIDR format (eg., 123.1.2.3/24) to be allowed.",
            nullable=True,
        )

        # define Arg Group "Properties"

        _args_schema = cls._args_schema
        _args_schema.disable_local_auth = AAZBoolArg(
            options=["--disable-local-auth"],
            arg_group="Properties",
            help="When set to true, calls to the search service will not be permitted to utilize API keys for authentication. This cannot be set to true if 'dataPlaneAuthOptions' are defined.",
            nullable=True,
        )
        _args_schema.partition_count = AAZIntArg(
            options=["--partition-count"],
            arg_group="Properties",
            help="The number of partitions in the search service; if specified, it can be 1, 2, 3, 4, 6, or 12. Values greater than 1 are only valid for standard SKUs. For 'standard3' services with hostingMode set to 'highDensity', the allowed values are between 1 and 3.",
            nullable=True,
            fmt=AAZIntArgFormat(
                maximum=12,
                minimum=1,
            ),
        )
        _args_schema.public_network_access = AAZStrArg(
            options=["--public-access", "--public-network-access"],
            arg_group="Properties",
            help="This value can be set to 'enabled' to avoid breaking changes on existing customer resources and templates. If set to 'disabled', traffic over public interface is not allowed, and private endpoint connections would be the exclusive access method.",
            nullable=True,
            enum={"disabled": "disabled", "enabled": "enabled"},
        )
        _args_schema.replica_count = AAZIntArg(
            options=["--replica-count"],
            arg_group="Properties",
            help="The number of replicas in the search service. If specified, it must be a value between 1 and 12 inclusive for standard SKUs or between 1 and 3 inclusive for basic SKU.",
            nullable=True,
            fmt=AAZIntArgFormat(
                maximum=12,
                minimum=1,
            ),
        )
        _args_schema.semantic_search = AAZStrArg(
            options=["--semantic-search"],
            arg_group="Properties",
            help="Sets options that control the availability of semantic search. This configuration is only possible for certain Azure Cognitive Search SKUs in certain locations.",
            nullable=True,
            enum={"disabled": "disabled", "free": "free", "standard": "standard"},
        )

        # define Arg Group "Service"

        _args_schema = cls._args_schema
        _args_schema.tags = AAZDictArg(
            options=["--tags"],
            arg_group="Service",
            help="Resource tags.",
            nullable=True,
        )

        tags = cls._args_schema.tags
        tags.Element = AAZStrArg(
            nullable=True,
        )

        # define Arg Group "Sku"
        return cls._args_schema

    def _execute_operations(self):
        self.pre_operations()
        self.ServicesGet(ctx=self.ctx)()
        self.pre_instance_update(self.ctx.vars.instance)
        self.InstanceUpdateByJson(ctx=self.ctx)()
        self.InstanceUpdateByGeneric(ctx=self.ctx)()
        self.post_instance_update(self.ctx.vars.instance)
        yield self.ServicesCreateOrUpdate(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 ServicesGet(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.Search/searchServices/{searchServiceName}",
                **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(
                    "resourceGroupName", self.ctx.args.resource_group,
                    required=True,
                ),
                **self.serialize_url_param(
                    "searchServiceName", self.ctx.args.search_service_name,
                    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", "2022-09-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_search_service_read(cls._schema_on_200)

            return cls._schema_on_200

    class ServicesCreateOrUpdate(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.Search/searchServices/{searchServiceName}",
                **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(
                    "resourceGroupName", self.ctx.args.resource_group,
                    required=True,
                ),
                **self.serialize_url_param(
                    "searchServiceName", self.ctx.args.search_service_name,
                    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", "2022-09-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_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()
            _UpdateHelper._build_schema_search_service_read(cls._schema_on_200_201)

            return cls._schema_on_200_201

    class InstanceUpdateByJson(AAZJsonInstanceUpdateOperation):

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

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

            identity = _builder.get(".identity")
            if identity is not None:
                identity.set_prop("type", AAZStrType, ".identity_type", typ_kwargs={"flags": {"required": True}})

            properties = _builder.get(".properties")
            if properties is not None:
                properties.set_prop("authOptions", AAZObjectType)
                properties.set_prop("disableLocalAuth", AAZBoolType, ".disable_local_auth", typ_kwargs={"nullable": True})
                properties.set_prop("networkRuleSet", AAZObjectType)
                properties.set_prop("partitionCount", AAZIntType, ".partition_count")
                properties.set_prop("publicNetworkAccess", AAZStrType, ".public_network_access")
                properties.set_prop("replicaCount", AAZIntType, ".replica_count")
                properties.set_prop("semanticSearch", AAZStrType, ".semantic_search", typ_kwargs={"nullable": True})

            auth_options = _builder.get(".properties.authOptions")
            if auth_options is not None:
                auth_options.set_prop("aadOrApiKey", AAZObjectType)
                auth_options.set_prop("apiKeyOnly", AAZObjectType, ".api_key_only")

            aad_or_api_key = _builder.get(".properties.authOptions.aadOrApiKey")
            if aad_or_api_key is not None:
                aad_or_api_key.set_prop("aadAuthFailureMode", AAZStrType, ".aad_auth_failure_mode")

            network_rule_set = _builder.get(".properties.networkRuleSet")
            if network_rule_set is not None:
                network_rule_set.set_prop("ipRules", AAZListType, ".ip_rules_internal")

            ip_rules = _builder.get(".properties.networkRuleSet.ipRules")
            if ip_rules is not None:
                ip_rules.set_elements(AAZObjectType, ".")

            _elements = _builder.get(".properties.networkRuleSet.ipRules[]")
            if _elements is not None:
                _elements.set_prop("value", AAZStrType, ".value")

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

            return _instance_value

    class InstanceUpdateByGeneric(AAZGenericInstanceUpdateOperation):

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


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

    _schema_search_service_read = None

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

        cls._schema_search_service_read = _schema_search_service_read = AAZObjectType()

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

        identity = _schema_search_service_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(
            flags={"required": True},
        )

        properties = _schema_search_service_read.properties
        properties.auth_options = AAZObjectType(
            serialized_name="authOptions",
        )
        properties.disable_local_auth = AAZBoolType(
            serialized_name="disableLocalAuth",
            nullable=True,
        )
        properties.encryption_with_cmk = AAZObjectType(
            serialized_name="encryptionWithCmk",
        )
        properties.hosting_mode = AAZStrType(
            serialized_name="hostingMode",
        )
        properties.network_rule_set = AAZObjectType(
            serialized_name="networkRuleSet",
        )
        properties.partition_count = AAZIntType(
            serialized_name="partitionCount",
        )
        properties.private_endpoint_connections = AAZListType(
            serialized_name="privateEndpointConnections",
            flags={"read_only": True},
        )
        properties.provisioning_state = AAZStrType(
            serialized_name="provisioningState",
            flags={"read_only": True},
        )
        properties.public_network_access = AAZStrType(
            serialized_name="publicNetworkAccess",
        )
        properties.replica_count = AAZIntType(
            serialized_name="replicaCount",
        )
        properties.semantic_search = AAZStrType(
            serialized_name="semanticSearch",
            nullable=True,
        )
        properties.shared_private_link_resources = AAZListType(
            serialized_name="sharedPrivateLinkResources",
            flags={"read_only": True},
        )
        properties.status = AAZStrType(
            flags={"read_only": True},
        )
        properties.status_details = AAZStrType(
            serialized_name="statusDetails",
            flags={"read_only": True},
        )

        auth_options = _schema_search_service_read.properties.auth_options
        auth_options.aad_or_api_key = AAZObjectType(
            serialized_name="aadOrApiKey",
        )
        auth_options.api_key_only = AAZObjectType(
            serialized_name="apiKeyOnly",
        )

        aad_or_api_key = _schema_search_service_read.properties.auth_options.aad_or_api_key
        aad_or_api_key.aad_auth_failure_mode = AAZStrType(
            serialized_name="aadAuthFailureMode",
        )

        encryption_with_cmk = _schema_search_service_read.properties.encryption_with_cmk
        encryption_with_cmk.encryption_compliance_status = AAZStrType(
            serialized_name="encryptionComplianceStatus",
            flags={"read_only": True},
        )
        encryption_with_cmk.enforcement = AAZStrType()

        network_rule_set = _schema_search_service_read.properties.network_rule_set
        network_rule_set.ip_rules = AAZListType(
            serialized_name="ipRules",
        )

        ip_rules = _schema_search_service_read.properties.network_rule_set.ip_rules
        ip_rules.Element = AAZObjectType()

        _element = _schema_search_service_read.properties.network_rule_set.ip_rules.Element
        _element.value = AAZStrType()

        private_endpoint_connections = _schema_search_service_read.properties.private_endpoint_connections
        private_endpoint_connections.Element = AAZObjectType()

        _element = _schema_search_service_read.properties.private_endpoint_connections.Element
        _element.id = AAZStrType(
            flags={"read_only": True},
        )
        _element.name = AAZStrType(
            flags={"read_only": True},
        )
        _element.properties = AAZObjectType()
        _element.type = AAZStrType(
            flags={"read_only": True},
        )

        properties = _schema_search_service_read.properties.private_endpoint_connections.Element.properties
        properties.group_id = AAZStrType(
            serialized_name="groupId",
        )
        properties.private_endpoint = AAZObjectType(
            serialized_name="privateEndpoint",
        )
        properties.private_link_service_connection_state = AAZObjectType(
            serialized_name="privateLinkServiceConnectionState",
        )
        properties.provisioning_state = AAZStrType(
            serialized_name="provisioningState",
        )

        private_endpoint = _schema_search_service_read.properties.private_endpoint_connections.Element.properties.private_endpoint
        private_endpoint.id = AAZStrType()

        private_link_service_connection_state = _schema_search_service_read.properties.private_endpoint_connections.Element.properties.private_link_service_connection_state
        private_link_service_connection_state.actions_required = AAZStrType(
            serialized_name="actionsRequired",
        )
        private_link_service_connection_state.description = AAZStrType()
        private_link_service_connection_state.status = AAZStrType()

        shared_private_link_resources = _schema_search_service_read.properties.shared_private_link_resources
        shared_private_link_resources.Element = AAZObjectType()

        _element = _schema_search_service_read.properties.shared_private_link_resources.Element
        _element.id = AAZStrType(
            flags={"read_only": True},
        )
        _element.name = AAZStrType(
            flags={"read_only": True},
        )
        _element.properties = AAZObjectType()
        _element.type = AAZStrType(
            flags={"read_only": True},
        )

        properties = _schema_search_service_read.properties.shared_private_link_resources.Element.properties
        properties.group_id = AAZStrType(
            serialized_name="groupId",
        )
        properties.private_link_resource_id = AAZStrType(
            serialized_name="privateLinkResourceId",
        )
        properties.provisioning_state = AAZStrType(
            serialized_name="provisioningState",
        )
        properties.request_message = AAZStrType(
            serialized_name="requestMessage",
        )
        properties.resource_region = AAZStrType(
            serialized_name="resourceRegion",
        )
        properties.status = AAZStrType()

        sku = _schema_search_service_read.sku
        sku.name = AAZStrType(
            flags={"required": True},
        )

        tags = _schema_search_service_read.tags
        tags.Element = AAZStrType()

        _schema.id = cls._schema_search_service_read.id
        _schema.identity = cls._schema_search_service_read.identity
        _schema.location = cls._schema_search_service_read.location
        _schema.name = cls._schema_search_service_read.name
        _schema.properties = cls._schema_search_service_read.properties
        _schema.sku = cls._schema_search_service_read.sku
        _schema.tags = cls._schema_search_service_read.tags
        _schema.type = cls._schema_search_service_read.type


__all__ = ["Update"]
