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

from typing import Any, Union, Optional, Dict
from marshmallow import fields, post_load
from azure.ml._schema.schema import PatchedNested
from .scale_settings_schema import ScaleSettingsSchema
from .request_settings_schema import InternalRequestSettings, RequestSettingsSchema
from .resource_requirements_schema import ResourceRequirementsSchema
from azure.ml._schema._endpoint.code_configuration_schema import InternalCodeConfiguration
from .liveness_probe import LivenessProbeSchema
from azure.ml._workspace_dependent_operations import WorkspaceScope
from azure.ml._utils._arm_id_utils import get_arm_id_with_version
from azure.ml._schema.environment import InternalEnvironment
from azure.ml._schema.model import InternalModel
from azure.ml._schema._endpoint.endpoint_deployment import EndpointDeploymentSchema, InternalEndpointDeployment
from azure.ml._restclient.machinelearningservices.models import (
    CodeConfiguration,
    DeploymentConfigurationBase,
    OnlineDeploymentPropertiesTrackedResource,
    OnlineDeploymentProperties,
    AksDeploymentConfiguration,
    ManagedDeploymentConfiguration,
    ContainerResourceRequirementsAutoGenerated,
    IdAssetReference,
    LivenessProbeRequirements,
    ScaleSettings,
)
from azure.ml.constants import BASE_PATH_CONTEXT_KEY, ComputeType


class InternalOnlineEndpointDeployment(InternalEndpointDeployment):
    def __init__(
        self,
        base_path: Optional[str] = None,
        id: str = None,
        name: str = None,
        type: str = None,
        tags: Dict[str, Any] = None,
        properties: Dict[str, Any] = None,
        model: Union[str, InternalModel] = None,
        code_configuration: InternalCodeConfiguration = None,
        environment: Union[str, InternalEnvironment] = None,
        sku: str = None,
        app_insights_enabled: bool = False,
        resource_requirements: ContainerResourceRequirementsAutoGenerated = None,
        scale_settings: ScaleSettings = None,
        request_settings: InternalRequestSettings = None,
        liveness_probe: LivenessProbeRequirements = None,
        provisioning_status: str = None,
    ):
        super(InternalOnlineEndpointDeployment, self).__init__(
            base_path=base_path,
            id=id,
            name=name,
            type=type,
            tags=tags,
            properties=properties,
            model=model,
            code_configuration=code_configuration,
            environment=environment,
        )

        self.sku = sku
        self.app_insights_enabled = app_insights_enabled
        self.resource_requirements = resource_requirements
        self.scale_settings = scale_settings
        self.request_settings = request_settings
        self.liveness_probe = liveness_probe
        self.provisioning_status = provisioning_status

    def _to_rest_deployment_configuration(self, compute_type: str) -> DeploymentConfigurationBase:
        concurrent_requests = (
            self.request_settings.max_concurrent_requests_per_instance if self.request_settings else None
        )
        if compute_type == ComputeType.AKS:
            return AksDeploymentConfiguration(
                app_insights_enabled=self.app_insights_enabled,
                max_queue_wait_ms=self.request_settings.max_queue_wait_ms if self.request_settings else None,
                scoring_timeout_ms=self.request_settings.request_timeout_ms if self.request_settings else None,
                max_concurrent_requests_per_instance=concurrent_requests,
                liveness_probe_requirements=self.liveness_probe,
                container_resource_requirements=self.resource_requirements,
            )
        elif compute_type == ComputeType.MANAGED:
            return ManagedDeploymentConfiguration(
                app_insights_enabled=self.app_insights_enabled,
                max_queue_wait_ms=self.request_settings.max_queue_wait_ms if self.request_settings else None,
                scoring_timeout_ms=self.request_settings.request_timeout_ms if self.request_settings else None,
                max_concurrent_requests_per_instance=concurrent_requests,
                liveness_probe_requirements=self.liveness_probe,
                instance_type=self.sku,
            )

    def _to_rest_online_deployments(
        self, compute_type: str, location: str, workspace_scope: WorkspaceScope
    ) -> OnlineDeploymentPropertiesTrackedResource:
        command = None
        if self.code_configuration.scoring_script:
            command = self.code_configuration.scoring_script
        if isinstance(self.code_configuration.code, str):
            code = CodeConfiguration(code_artifact_id=self.code_configuration.code, command=command)
        else:
            code = CodeConfiguration(
                code_artifact_id=get_arm_id_with_version(
                    workspace_scope, "codes", self.code_configuration.code.name, self.code_configuration.code.version
                ),
                command=command,
            )
        if isinstance(self.model, str):
            model = IdAssetReference(id=self.model)
        else:
            model = IdAssetReference(
                id=get_arm_id_with_version(workspace_scope, "models", self.model.name, self.model.version)
            )
        if isinstance(self.environment, str):
            environment = self.environment
        else:
            environment = get_arm_id_with_version(
                workspace_scope, "environments", self.environment.name, self.environment.version
            )
        deployment_config = self._to_rest_deployment_configuration(compute_type=compute_type)
        properties = OnlineDeploymentProperties(
            code_configuration=code,
            environment_id=environment,
            model_reference=model,
            deployment_configuration=deployment_config,
            scale_settings=self.scale_settings,
            properties=self.properties,
        )
        return OnlineDeploymentPropertiesTrackedResource(location=location, properties=properties, tags=self.tags)

    def _from_rest_online_deployment(self, deployment: OnlineDeploymentPropertiesTrackedResource):
        liveness_probe = None
        request_settings = None
        deployment_config = deployment.properties.deployment_configuration
        if isinstance(deployment_config, AksDeploymentConfiguration):
            liveness_probe = deployment.properties.deployment_configuration.liveness_probe_requirements
            request_settings = InternalRequestSettings(
                request_timeout_ms=deployment_config.scoring_timeout_ms,
                max_concurrent_requests_per_instance=deployment_config.max_concurrent_requests_per_instance,
                max_queue_wait_ms=deployment_config.max_queue_wait_ms,
            )
        self.id = deployment.id
        self.name = deployment.name
        self.type = deployment.type
        self.tags = deployment.tags
        self.properties = deployment.properties.properties
        self.model = deployment.properties.model_reference.id
        self.code_configuration = InternalCodeConfiguration(
            code=deployment.properties.code_configuration.code_artifact_id,
            scoring_script=deployment.properties.code_configuration.command[0],
        )
        self.environment = deployment.properties.environment_id
        self.sku = None
        self.app_insights_enabled = deployment.properties.deployment_configuration.app_insights_enabled
        self.resource_requirements = deployment.properties.deployment_configuration.container_resource_requirements
        self.scale_settings = deployment.properties.scale_settings
        self.request_settings = request_settings
        self.liveness_probe = liveness_probe
        self.provisioning_status = deployment.properties.provisioning_state
        return self


class OnlineEndpointDeploymentSchema(EndpointDeploymentSchema):
    sku = fields.Str()
    app_insights_enabled = fields.Bool()
    resource_requirements = PatchedNested(ResourceRequirementsSchema, required=False)
    scale_settings = PatchedNested(ScaleSettingsSchema)
    request_settings = PatchedNested(RequestSettingsSchema)
    liveness_probe = PatchedNested(LivenessProbeSchema)
    provisioning_status = fields.Str(dump_only=True)

    @post_load
    def make(self, data: Any, **kwargs: Any) -> InternalOnlineEndpointDeployment:
        return InternalOnlineEndpointDeployment(base_path=self.context[BASE_PATH_CONTEXT_KEY], **data)


class NamedOnlineEndpointDeploymentSchema(OnlineEndpointDeploymentSchema):
    name = fields.Str(required=True)
