# --------------------------------------------------------------------------------------------
# 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(
    "network watcher test-connectivity",
    is_preview=True,
)
class TestConnectivity(AAZCommand):
    """Test if a connection can be established between a Virtual Machine and a given endpoint.

    To check connectivity between two VMs in different regions, use the VM ids instead of the VM names for the source and destination resource arguments. To register for this feature or see additional examples visit https://learn.microsoft.com/en-us/azure/network-watcher/network-watcher-connectivity-cli.

    :example: Check connectivity between two virtual machines in the same resource group over port 80.
        az network watcher test-connectivity -g MyResourceGroup --source-resource MyVmName1 --dest-resource MyVmName2 --dest-port 80

    :example: Check connectivity between two virtual machines in the same subscription in two different resource groups over port 80.
        az network watcher test-connectivity --source-resource MyVmId1 --dest-resource MyVmId2 --dest-port 80
    """

    _aaz_info = {
        "version": "2022-01-01",
        "resources": [
            ["mgmt-plane", "/subscriptions/{}/resourcegroups/{}/providers/microsoft.network/networkwatchers/{}/connectivitycheck", "2022-01-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.watcher_name = AAZStrArg(
            options=["--watcher-name"],
            help="Name of the network watcher resource.",
            required=True,
            id_part="name",
        )
        _args_schema.watcher_rg = AAZResourceGroupNameArg(
            options=["--watcher-rg"],
            help="Name of the resource group the watcher is in.",
            required=True,
        )
        _args_schema.protocol = AAZStrArg(
            options=["--protocol"],
            help="Protocol to test on.",
            enum={"Http": "Http", "Https": "Https", "Icmp": "Icmp", "Tcp": "Tcp"},
        )
        _args_schema.source_port = AAZIntArg(
            options=["--source-port"],
            help="Port number from which to originate traffic.",
            fmt=AAZIntArgFormat(
                maximum=65535,
                minimum=0,
            ),
        )
        _args_schema.source_resource = AAZStrArg(
            options=["--source-resource"],
            help="Name or ID of the resource from which to originate traffic. Currently only Virtual Machines are supported.",
            required=True,
        )

        # define Arg Group "Destination"

        _args_schema = cls._args_schema
        _args_schema.dest_address = AAZStrArg(
            options=["--dest-address"],
            arg_group="Destination",
            help="IP address or URI at which to receive traffic.",
        )
        _args_schema.dest_port = AAZIntArg(
            options=["--dest-port"],
            arg_group="Destination",
            help="Port number on which to receive traffic.",
            fmt=AAZIntArgFormat(
                maximum=65535,
                minimum=0,
            ),
        )
        _args_schema.dest_resource = AAZStrArg(
            options=["--dest-resource"],
            arg_group="Destination",
            help="Name or ID of the resource to receive traffic. Currently only Virtual Machines are supported.",
        )

        # define Arg Group "HTTP Configuration"

        _args_schema = cls._args_schema
        _args_schema.headers_obj = AAZListArg(
            options=["--headers-obj"],
            arg_group="HTTP Configuration",
            help="List of HTTP headers.",
        )
        _args_schema.method = AAZStrArg(
            options=["--method"],
            arg_group="HTTP Configuration",
            help="HTTP method to use.",
            enum={"Get": "Get"},
        )
        _args_schema.valid_status_codes = AAZListArg(
            options=["--valid-status-codes"],
            arg_group="HTTP Configuration",
            help="Space-separated list of HTTP status codes considered valid.",
        )

        headers_obj = cls._args_schema.headers_obj
        headers_obj.Element = AAZObjectArg()

        _element = cls._args_schema.headers_obj.Element
        _element.name = AAZStrArg(
            options=["name"],
            help="The name in HTTP header.",
        )
        _element.value = AAZStrArg(
            options=["value"],
            help="The value in HTTP header.",
        )

        valid_status_codes = cls._args_schema.valid_status_codes
        valid_status_codes.Element = AAZIntArg()

        # define Arg Group "Parameters"
        return cls._args_schema

    def _execute_operations(self):
        self.pre_operations()
        yield self.NetworkWatchersCheckConnectivity(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 NetworkWatchersCheckConnectivity(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": "location"},
                    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": "location"},
                    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.Network/networkWatchers/{networkWatcherName}/connectivityCheck",
                **self.url_parameters
            )

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

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

        @property
        def url_parameters(self):
            parameters = {
                **self.serialize_url_param(
                    "networkWatcherName", self.ctx.args.watcher_name,
                    required=True,
                ),
                **self.serialize_url_param(
                    "resourceGroupName", self.ctx.args.watcher_rg,
                    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-01-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("destination", AAZObjectType, ".", typ_kwargs={"flags": {"required": True}})
            _builder.set_prop("protocol", AAZStrType, ".protocol")
            _builder.set_prop("protocolConfiguration", AAZObjectType)
            _builder.set_prop("source", AAZObjectType, ".", typ_kwargs={"flags": {"required": True}})

            destination = _builder.get(".destination")
            if destination is not None:
                destination.set_prop("address", AAZStrType, ".dest_address")
                destination.set_prop("port", AAZIntType, ".dest_port")
                destination.set_prop("resourceId", AAZStrType, ".dest_resource")

            protocol_configuration = _builder.get(".protocolConfiguration")
            if protocol_configuration is not None:
                protocol_configuration.set_prop("HTTPConfiguration", AAZObjectType)

            http_configuration = _builder.get(".protocolConfiguration.HTTPConfiguration")
            if http_configuration is not None:
                http_configuration.set_prop("headers", AAZListType, ".headers_obj")
                http_configuration.set_prop("method", AAZStrType, ".method")
                http_configuration.set_prop("validStatusCodes", AAZListType, ".valid_status_codes")

            headers = _builder.get(".protocolConfiguration.HTTPConfiguration.headers")
            if headers is not None:
                headers.set_elements(AAZObjectType, ".")

            _elements = _builder.get(".protocolConfiguration.HTTPConfiguration.headers[]")
            if _elements is not None:
                _elements.set_prop("name", AAZStrType, ".name")
                _elements.set_prop("value", AAZStrType, ".value")

            valid_status_codes = _builder.get(".protocolConfiguration.HTTPConfiguration.validStatusCodes")
            if valid_status_codes is not None:
                valid_status_codes.set_elements(AAZIntType, ".")

            source = _builder.get(".source")
            if source is not None:
                source.set_prop("port", AAZIntType, ".source_port")
                source.set_prop("resourceId", AAZStrType, ".source_resource", typ_kwargs={"flags": {"required": True}})

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

            return cls._schema_on_200


class _TestConnectivityHelper:
    """Helper class for TestConnectivity"""

    _schema_connectivity_information_read = None

    @classmethod
    def _build_schema_connectivity_information_read(cls, _schema):
        if cls._schema_connectivity_information_read is not None:
            _schema.avg_latency_in_ms = cls._schema_connectivity_information_read.avg_latency_in_ms
            _schema.connection_status = cls._schema_connectivity_information_read.connection_status
            _schema.hops = cls._schema_connectivity_information_read.hops
            _schema.max_latency_in_ms = cls._schema_connectivity_information_read.max_latency_in_ms
            _schema.min_latency_in_ms = cls._schema_connectivity_information_read.min_latency_in_ms
            _schema.probes_failed = cls._schema_connectivity_information_read.probes_failed
            _schema.probes_sent = cls._schema_connectivity_information_read.probes_sent
            return

        cls._schema_connectivity_information_read = _schema_connectivity_information_read = AAZObjectType()

        connectivity_information_read = _schema_connectivity_information_read
        connectivity_information_read.avg_latency_in_ms = AAZIntType(
            serialized_name="avgLatencyInMs",
            flags={"read_only": True},
        )
        connectivity_information_read.connection_status = AAZStrType(
            serialized_name="connectionStatus",
            flags={"read_only": True},
        )
        connectivity_information_read.hops = AAZListType(
            flags={"read_only": True},
        )
        connectivity_information_read.max_latency_in_ms = AAZIntType(
            serialized_name="maxLatencyInMs",
            flags={"read_only": True},
        )
        connectivity_information_read.min_latency_in_ms = AAZIntType(
            serialized_name="minLatencyInMs",
            flags={"read_only": True},
        )
        connectivity_information_read.probes_failed = AAZIntType(
            serialized_name="probesFailed",
            flags={"read_only": True},
        )
        connectivity_information_read.probes_sent = AAZIntType(
            serialized_name="probesSent",
            flags={"read_only": True},
        )

        hops = _schema_connectivity_information_read.hops
        hops.Element = AAZObjectType()

        _element = _schema_connectivity_information_read.hops.Element
        _element.address = AAZStrType(
            flags={"read_only": True},
        )
        _element.id = AAZStrType(
            flags={"read_only": True},
        )
        _element.issues = AAZListType(
            flags={"read_only": True},
        )
        _element.links = AAZListType(
            flags={"read_only": True},
        )
        _element.next_hop_ids = AAZListType(
            serialized_name="nextHopIds",
            flags={"read_only": True},
        )
        _element.previous_hop_ids = AAZListType(
            serialized_name="previousHopIds",
            flags={"read_only": True},
        )
        _element.previous_links = AAZListType(
            serialized_name="previousLinks",
            flags={"read_only": True},
        )
        _element.resource_id = AAZStrType(
            serialized_name="resourceId",
            flags={"read_only": True},
        )
        _element.type = AAZStrType(
            flags={"read_only": True},
        )

        issues = _schema_connectivity_information_read.hops.Element.issues
        issues.Element = AAZObjectType()
        cls._build_schema_connectivity_issue_read(issues.Element)

        links = _schema_connectivity_information_read.hops.Element.links
        links.Element = AAZObjectType()
        cls._build_schema_hop_link_read(links.Element)

        next_hop_ids = _schema_connectivity_information_read.hops.Element.next_hop_ids
        next_hop_ids.Element = AAZStrType()

        previous_hop_ids = _schema_connectivity_information_read.hops.Element.previous_hop_ids
        previous_hop_ids.Element = AAZStrType()

        previous_links = _schema_connectivity_information_read.hops.Element.previous_links
        previous_links.Element = AAZObjectType()
        cls._build_schema_hop_link_read(previous_links.Element)

        _schema.avg_latency_in_ms = cls._schema_connectivity_information_read.avg_latency_in_ms
        _schema.connection_status = cls._schema_connectivity_information_read.connection_status
        _schema.hops = cls._schema_connectivity_information_read.hops
        _schema.max_latency_in_ms = cls._schema_connectivity_information_read.max_latency_in_ms
        _schema.min_latency_in_ms = cls._schema_connectivity_information_read.min_latency_in_ms
        _schema.probes_failed = cls._schema_connectivity_information_read.probes_failed
        _schema.probes_sent = cls._schema_connectivity_information_read.probes_sent

    _schema_connectivity_issue_read = None

    @classmethod
    def _build_schema_connectivity_issue_read(cls, _schema):
        if cls._schema_connectivity_issue_read is not None:
            _schema.context = cls._schema_connectivity_issue_read.context
            _schema.origin = cls._schema_connectivity_issue_read.origin
            _schema.severity = cls._schema_connectivity_issue_read.severity
            _schema.type = cls._schema_connectivity_issue_read.type
            return

        cls._schema_connectivity_issue_read = _schema_connectivity_issue_read = AAZObjectType()

        connectivity_issue_read = _schema_connectivity_issue_read
        connectivity_issue_read.context = AAZListType(
            flags={"read_only": True},
        )
        connectivity_issue_read.origin = AAZStrType(
            flags={"read_only": True},
        )
        connectivity_issue_read.severity = AAZStrType(
            flags={"read_only": True},
        )
        connectivity_issue_read.type = AAZStrType(
            flags={"read_only": True},
        )

        context = _schema_connectivity_issue_read.context
        context.Element = AAZDictType()

        _element = _schema_connectivity_issue_read.context.Element
        _element.Element = AAZStrType()

        _schema.context = cls._schema_connectivity_issue_read.context
        _schema.origin = cls._schema_connectivity_issue_read.origin
        _schema.severity = cls._schema_connectivity_issue_read.severity
        _schema.type = cls._schema_connectivity_issue_read.type

    _schema_hop_link_read = None

    @classmethod
    def _build_schema_hop_link_read(cls, _schema):
        if cls._schema_hop_link_read is not None:
            _schema.context = cls._schema_hop_link_read.context
            _schema.issues = cls._schema_hop_link_read.issues
            _schema.link_type = cls._schema_hop_link_read.link_type
            _schema.next_hop_id = cls._schema_hop_link_read.next_hop_id
            _schema.properties = cls._schema_hop_link_read.properties
            _schema.resource_id = cls._schema_hop_link_read.resource_id
            return

        cls._schema_hop_link_read = _schema_hop_link_read = AAZObjectType()

        hop_link_read = _schema_hop_link_read
        hop_link_read.context = AAZDictType(
            flags={"read_only": True},
        )
        hop_link_read.issues = AAZListType(
            flags={"read_only": True},
        )
        hop_link_read.link_type = AAZStrType(
            serialized_name="linkType",
            flags={"read_only": True},
        )
        hop_link_read.next_hop_id = AAZStrType(
            serialized_name="nextHopId",
            flags={"read_only": True},
        )
        hop_link_read.properties = AAZObjectType(
            flags={"client_flatten": True},
        )
        hop_link_read.resource_id = AAZStrType(
            serialized_name="resourceId",
            flags={"read_only": True},
        )

        context = _schema_hop_link_read.context
        context.Element = AAZStrType()

        issues = _schema_hop_link_read.issues
        issues.Element = AAZObjectType()
        cls._build_schema_connectivity_issue_read(issues.Element)

        properties = _schema_hop_link_read.properties
        properties.round_trip_time_avg = AAZIntType(
            serialized_name="roundTripTimeAvg",
            flags={"read_only": True},
        )
        properties.round_trip_time_max = AAZIntType(
            serialized_name="roundTripTimeMax",
            flags={"read_only": True},
        )
        properties.round_trip_time_min = AAZIntType(
            serialized_name="roundTripTimeMin",
            flags={"read_only": True},
        )

        _schema.context = cls._schema_hop_link_read.context
        _schema.issues = cls._schema_hop_link_read.issues
        _schema.link_type = cls._schema_hop_link_read.link_type
        _schema.next_hop_id = cls._schema_hop_link_read.next_hop_id
        _schema.properties = cls._schema_hop_link_read.properties
        _schema.resource_id = cls._schema_hop_link_read.resource_id


__all__ = ["TestConnectivity"]
