# ---------------------------------------------------------
# Copyright (c) Microsoft Corporation. All rights reserved.
# ---------------------------------------------------------

from abc import ABC
from typing import Any, Dict, List, Optional

from azure.ai.ml._restclient.v2024_10_01_preview.models import FqdnOutboundRule as RestFqdnOutboundRule
from azure.ai.ml._restclient.v2024_10_01_preview.models import (
    ManagedNetworkProvisionStatus as RestManagedNetworkProvisionStatus,
)
from azure.ai.ml._restclient.v2024_10_01_preview.models import ManagedNetworkSettings as RestManagedNetwork
from azure.ai.ml._restclient.v2024_10_01_preview.models import (
    PrivateEndpointDestination as RestPrivateEndpointOutboundRuleDestination,
)
from azure.ai.ml._restclient.v2024_10_01_preview.models import (
    PrivateEndpointOutboundRule as RestPrivateEndpointOutboundRule,
)
from azure.ai.ml._restclient.v2024_10_01_preview.models import (
    ServiceTagDestination as RestServiceTagOutboundRuleDestination,
)
from azure.ai.ml._restclient.v2024_10_01_preview.models import ServiceTagOutboundRule as RestServiceTagOutboundRule
from azure.ai.ml.constants._workspace import IsolationMode, OutboundRuleCategory, OutboundRuleType


class OutboundRule(ABC):
    """Base class for Outbound Rules, cannot be instantiated directly. Please see FqdnDestination,
    PrivateEndpointDestination, and ServiceTagDestination objects to create outbound rules.

    :param name: Name of the outbound rule.
    :type name: str
    :param type: Type of the outbound rule. Supported types are "FQDN", "PrivateEndpoint", "ServiceTag"
    :type type: str
    :ivar type: Type of the outbound rule. Supported types are "FQDN", "PrivateEndpoint", "ServiceTag"
    :vartype type: str
    """

    def __init__(
        self,
        *,
        name: Optional[str] = None,
        **kwargs: Any,
    ) -> None:
        self.name = name
        self.parent_rule_names = kwargs.pop("parent_rule_names", None)
        self.type = kwargs.pop("type", None)
        self.category = kwargs.pop("category", OutboundRuleCategory.USER_DEFINED)
        self.status = kwargs.pop("status", None)

    @classmethod
    def _from_rest_object(cls, rest_obj: Any, name: str) -> Optional["OutboundRule"]:
        if isinstance(rest_obj, RestFqdnOutboundRule):
            rule_fqdnDestination = FqdnDestination(destination=rest_obj.destination, name=name)
            rule_fqdnDestination.category = rest_obj.category
            rule_fqdnDestination.status = rest_obj.status
            return rule_fqdnDestination
        if isinstance(rest_obj, RestPrivateEndpointOutboundRule):
            rule_privateEndpointDestination = PrivateEndpointDestination(
                service_resource_id=rest_obj.destination.service_resource_id,
                subresource_target=rest_obj.destination.subresource_target,
                spark_enabled=rest_obj.destination.spark_enabled,
                fqdns=rest_obj.fqdns,
                name=name,
            )
            rule_privateEndpointDestination.category = rest_obj.category
            rule_privateEndpointDestination.status = rest_obj.status
            return rule_privateEndpointDestination
        if isinstance(rest_obj, RestServiceTagOutboundRule):
            rule = ServiceTagDestination(
                service_tag=rest_obj.destination.service_tag,
                protocol=rest_obj.destination.protocol,
                port_ranges=rest_obj.destination.port_ranges,
                address_prefixes=rest_obj.destination.address_prefixes,
                name=name,
            )
            rule.category = rest_obj.category
            rule.status = rest_obj.status
            return rule

        return None


class FqdnDestination(OutboundRule):
    """Class representing a FQDN outbound rule.

    :param name: Name of the outbound rule.
    :type name: str
    :param destination: Fully qualified domain name to which outbound connections are allowed.
        For example: “xxxxxx.contoso.com”.
    :type destination: str
    :ivar type: Type of the outbound rule. Set to "FQDN" for this class.
    :vartype type: str

    .. literalinclude:: ../samples/ml_samples_workspace.py
            :start-after: [START fqdn_outboundrule]
            :end-before: [END fqdn_outboundrule]
            :language: python
            :dedent: 8
            :caption: Creating a FqdnDestination outbound rule object.
    """

    def __init__(self, *, name: str, destination: str, **kwargs: Any) -> None:
        self.destination = destination
        OutboundRule.__init__(self, type=OutboundRuleType.FQDN, name=name, **kwargs)

    def _to_rest_object(self) -> RestFqdnOutboundRule:
        return RestFqdnOutboundRule(type=self.type, category=self.category, destination=self.destination)

    def _to_dict(self) -> Dict:
        return {
            "name": self.name,
            "type": OutboundRuleType.FQDN,
            "category": self.category,
            "destination": self.destination,
            "status": self.status,
        }


class PrivateEndpointDestination(OutboundRule):
    """Class representing a Private Endpoint outbound rule.

    :param name: Name of the outbound rule.
    :type name: str
    :param service_resource_id: The resource URI of the root service that supports creation of the private link.
    :type service_resource_id: str
    :param subresource_target: The target endpoint of the subresource of the service.
    :type subresource_target: str
    :param spark_enabled: Indicates if the private endpoint can be used for Spark jobs, default is “false”.
    :type spark_enabled: bool
    :param fqdns: String list of FQDNs particular to the Private Endpoint resource creation. For application
        gateway Private Endpoints, this is the FQDN which will resolve to the private IP of the application
        gateway PE inside the workspace's managed network.
    :type fqdns: List[str]
    :ivar type: Type of the outbound rule. Set to "PrivateEndpoint" for this class.
    :vartype type: str

    .. literalinclude:: ../samples/ml_samples_workspace.py
            :start-after: [START private_endpoint_outboundrule]
            :end-before: [END private_endpoint_outboundrule]
            :language: python
            :dedent: 8
            :caption: Creating a PrivateEndpointDestination outbound rule object.
    """

    def __init__(
        self,
        *,
        name: str,
        service_resource_id: str,
        subresource_target: str,
        spark_enabled: bool = False,
        fqdns: Optional[List[str]] = None,
        **kwargs: Any,
    ) -> None:
        self.service_resource_id = service_resource_id
        self.subresource_target = subresource_target
        self.spark_enabled = spark_enabled
        self.fqdns = fqdns
        OutboundRule.__init__(self, type=OutboundRuleType.PRIVATE_ENDPOINT, name=name, **kwargs)

    def _to_rest_object(self) -> RestPrivateEndpointOutboundRule:
        return RestPrivateEndpointOutboundRule(
            type=self.type,
            category=self.category,
            destination=RestPrivateEndpointOutboundRuleDestination(
                service_resource_id=self.service_resource_id,
                subresource_target=self.subresource_target,
                spark_enabled=self.spark_enabled,
            ),
            fqdns=self.fqdns,
        )

    def _to_dict(self) -> Dict:
        return {
            "name": self.name,
            "type": OutboundRuleType.PRIVATE_ENDPOINT,
            "category": self.category,
            "destination": {
                "service_resource_id": self.service_resource_id,
                "subresource_target": self.subresource_target,
                "spark_enabled": self.spark_enabled,
            },
            "fqdns": self.fqdns,
            "status": self.status,
        }


class ServiceTagDestination(OutboundRule):
    """Class representing a Service Tag outbound rule.

    :param name: Name of the outbound rule.
    :type name: str
    :param service_tag: Service Tag of an Azure service, maps to predefined IP addresses for its service endpoints.
    :type service_tag: str
    :param protocol: Allowed transport protocol, can be "TCP", "UDP", "ICMP" or "*" for all supported protocols.
    :type protocol: str
    :param port_ranges: A comma-separated list of single ports and/or range of ports, such as "80,1024-65535".
        Traffics should be allowed to these port ranges.
    :type port_ranges: str
    :param address_prefixes: Optional list of CIDR prefixes or IP ranges, when provided, service_tag argument will
        be ignored and address_prefixes will be used instead.
    :type address_prefixes: List[str]
    :ivar type: Type of the outbound rule. Set to "ServiceTag" for this class.
    :vartype type: str

    .. literalinclude:: ../samples/ml_samples_workspace.py
            :start-after: [START service_tag_outboundrule]
            :end-before: [END service_tag_outboundrule]
            :language: python
            :dedent: 8
            :caption: Creating a ServiceTagDestination outbound rule object.
    """

    def __init__(
        self,
        *,
        name: str,
        protocol: str,
        port_ranges: str,
        service_tag: Optional[str] = None,
        address_prefixes: Optional[List[str]] = None,
        **kwargs: Any,
    ) -> None:
        self.service_tag = service_tag
        self.protocol = protocol
        self.port_ranges = port_ranges
        self.address_prefixes = address_prefixes
        OutboundRule.__init__(self, type=OutboundRuleType.SERVICE_TAG, name=name, **kwargs)

    def _to_rest_object(self) -> RestServiceTagOutboundRule:
        return RestServiceTagOutboundRule(
            type=self.type,
            category=self.category,
            destination=RestServiceTagOutboundRuleDestination(
                service_tag=self.service_tag,
                protocol=self.protocol,
                port_ranges=self.port_ranges,
                address_prefixes=self.address_prefixes,
            ),
        )

    def _to_dict(self) -> Dict:
        return {
            "name": self.name,
            "type": OutboundRuleType.SERVICE_TAG,
            "category": self.category,
            "destination": {
                "service_tag": self.service_tag,
                "protocol": self.protocol,
                "port_ranges": self.port_ranges,
                "address_prefixes": self.address_prefixes,
            },
            "status": self.status,
        }


class ManagedNetwork:
    """Managed Network settings for a workspace.

    :param isolation_mode: Isolation of the managed network, defaults to Disabled.
    :type isolation_mode: str
    :param firewall_sku: Firewall Sku for FQDN rules in AllowOnlyApprovedOutbound..
    :type firewall_sku: str
    :param outbound_rules: List of outbound rules for the managed network.
    :type outbound_rules: List[~azure.ai.ml.entities.OutboundRule]
    :param network_id: Network id for the managed network, not meant to be set by user.
    :type network_id: str

    .. literalinclude:: ../samples/ml_samples_workspace.py
            :start-after: [START workspace_managed_network]
            :end-before: [END workspace_managed_network]
            :language: python
            :dedent: 8
            :caption: Creating a ManagedNetwork object with one of each rule type.
    """

    def __init__(
        self,
        *,
        isolation_mode: str = IsolationMode.DISABLED,
        outbound_rules: Optional[List[OutboundRule]] = None,
        firewall_sku: Optional[str] = None,
        network_id: Optional[str] = None,
        **kwargs: Any,
    ) -> None:
        self.isolation_mode = isolation_mode
        self.firewall_sku = firewall_sku
        self.network_id = network_id
        self.outbound_rules = outbound_rules
        self.status = kwargs.pop("status", None)

    def _to_rest_object(self) -> RestManagedNetwork:
        rest_outbound_rules = (
            {
                # pylint: disable=protected-access
                outbound_rule.name: outbound_rule._to_rest_object()  # type: ignore[attr-defined]
                for outbound_rule in self.outbound_rules
            }
            if self.outbound_rules
            else {}
        )
        return RestManagedNetwork(
            isolation_mode=self.isolation_mode, outbound_rules=rest_outbound_rules, firewall_sku=self.firewall_sku
        )

    @classmethod
    def _from_rest_object(cls, obj: RestManagedNetwork) -> "ManagedNetwork":
        from_rest_outbound_rules = (
            [
                OutboundRule._from_rest_object(obj.outbound_rules[name], name=name)  # pylint: disable=protected-access
                for name in obj.outbound_rules
            ]
            if obj.outbound_rules
            else {}
        )
        return ManagedNetwork(
            isolation_mode=obj.isolation_mode,
            outbound_rules=from_rest_outbound_rules,  # type: ignore[arg-type]
            network_id=obj.network_id,
            status=obj.status,
            firewall_sku=obj.firewall_sku,
        )


class ManagedNetworkProvisionStatus:
    """ManagedNetworkProvisionStatus.

    :param status: Status for managed network provision.
    :type status: str
    :param spark_ready: Bool value indicating if managed network is spark ready
    :type spark_ready: bool
    """

    def __init__(
        self,
        *,
        status: Optional[str] = None,
        spark_ready: Optional[bool] = None,
    ):
        self.status = status
        self.spark_ready = spark_ready

    @classmethod
    def _from_rest_object(cls, rest_obj: RestManagedNetworkProvisionStatus) -> "ManagedNetworkProvisionStatus":
        return cls(
            status=rest_obj.status,
            spark_ready=rest_obj.spark_ready,
        )
