# -*- coding: utf-8 -*-

# Copyright 2020 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#

import pytest

from unittest import mock
from importlib import reload
from datetime import datetime, timedelta

from google.api_core import operation as ga_operation
from google.auth import credentials as auth_credentials

from google.cloud import aiplatform

from google.cloud.aiplatform import initializer
from google.cloud.aiplatform import models
from google.cloud.aiplatform import utils

from google.cloud.aiplatform_v1beta1.services.endpoint_service import (
    client as endpoint_service_client_v1beta1,
)
from google.cloud.aiplatform_v1beta1.services.prediction_service import (
    client as prediction_service_client_v1beta1,
)
from google.cloud.aiplatform_v1beta1.types import (
    endpoint as gca_endpoint_v1beta1,
    machine_resources as gca_machine_resources_v1beta1,
    prediction_service as gca_prediction_service_v1beta1,
    endpoint_service as gca_endpoint_service_v1beta1,
)

from google.cloud.aiplatform_v1.services.model_service import (
    client as model_service_client,
)
from google.cloud.aiplatform_v1.services.endpoint_service import (
    client as endpoint_service_client,
)
from google.cloud.aiplatform_v1.services.prediction_service import (
    client as prediction_service_client,
)
from google.cloud.aiplatform_v1.types import (
    endpoint as gca_endpoint,
    model as gca_model,
    machine_resources as gca_machine_resources,
    prediction_service as gca_prediction_service,
    endpoint_service as gca_endpoint_service,
    encryption_spec as gca_encryption_spec,
)

_TEST_PROJECT = "test-project"
_TEST_PROJECT_2 = "test-project-2"
_TEST_LOCATION = "us-central1"
_TEST_LOCATION_2 = "europe-west4"

_TEST_DISPLAY_NAME = "test-display-name"
_TEST_DISPLAY_NAME_2 = "test-display-name-2"
_TEST_ID = "1028944691210842416"
_TEST_ID_2 = "4366591682456584192"
_TEST_DESCRIPTION = "test-description"

_TEST_ENDPOINT_NAME = (
    f"projects/{_TEST_PROJECT}/locations/{_TEST_LOCATION}/endpoints/{_TEST_ID}"
)
_TEST_ENDPOINT_NAME_ALT_LOCATION = (
    f"projects/{_TEST_PROJECT}/locations/{_TEST_LOCATION_2}/endpoints/{_TEST_ID}"
)
_TEST_PARENT = f"projects/{_TEST_PROJECT}/locations/{_TEST_LOCATION}"
_TEST_MODEL_NAME = (
    f"projects/{_TEST_PROJECT}/locations/{_TEST_LOCATION}/models/{_TEST_ID}"
)

_TEST_MODEL_ID = "1028944691210842416"
_TEST_PREDICTION = [[1.0, 2.0, 3.0], [3.0, 3.0, 1.0]]
_TEST_INSTANCES = [[1.0, 2.0, 3.0], [1.0, 3.0, 4.0]]
_TEST_CREDENTIALS = mock.Mock(spec=auth_credentials.AnonymousCredentials())
_TEST_SERVICE_ACCOUNT = "vinnys@my-project.iam.gserviceaccount.com"

_TEST_DEPLOYED_MODELS = [
    gca_endpoint.DeployedModel(id=_TEST_ID, display_name=_TEST_DISPLAY_NAME),
    gca_endpoint.DeployedModel(id=_TEST_ID_2, display_name=_TEST_DISPLAY_NAME_2),
]

_TEST_MACHINE_TYPE = "n1-standard-32"
_TEST_ACCELERATOR_TYPE = "NVIDIA_TESLA_P100"
_TEST_ACCELERATOR_COUNT = 2

_TEST_EXPLANATIONS = [
    gca_prediction_service_v1beta1.explanation.Explanation(attributions=[])
]

_TEST_ATTRIBUTIONS = [
    gca_prediction_service_v1beta1.explanation.Attribution(
        baseline_output_value=1.0,
        instance_output_value=2.0,
        feature_attributions=3.0,
        output_index=[1, 2, 3],
        output_display_name="abc",
        approximation_error=6.0,
        output_name="xyz",
    )
]

_TEST_EXPLANATION_METADATA = aiplatform.explain.ExplanationMetadata(
    inputs={
        "features": aiplatform.explain.ExplanationMetadata.InputMetadata(
            {
                "input_tensor_name": "dense_input",
                "encoding": "BAG_OF_FEATURES",
                "modality": "numeric",
                "index_feature_mapping": ["abc", "def", "ghj"],
            }
        )
    },
    outputs={
        "medv": aiplatform.explain.ExplanationMetadata.OutputMetadata(
            {"output_tensor_name": "dense_2"}
        )
    },
)
_TEST_EXPLANATION_PARAMETERS = aiplatform.explain.ExplanationParameters(
    {"sampled_shapley_attribution": {"path_count": 10}}
)

# CMEK encryption
_TEST_ENCRYPTION_KEY_NAME = "key_1234"
_TEST_ENCRYPTION_SPEC = gca_encryption_spec.EncryptionSpec(
    kms_key_name=_TEST_ENCRYPTION_KEY_NAME
)

_TEST_ENDPOINT_GAPIC = gca_endpoint.Endpoint(
    display_name=_TEST_DISPLAY_NAME, name=_TEST_ENDPOINT_NAME
)

_TEST_ENDPOINT_LIST = [
    gca_endpoint.Endpoint(
        name=_TEST_ENDPOINT_NAME,
        display_name="aac",
        create_time=datetime.now() - timedelta(minutes=15),
    ),
    gca_endpoint.Endpoint(
        name=_TEST_ENDPOINT_NAME,
        display_name="aab",
        create_time=datetime.now() - timedelta(minutes=5),
    ),
    gca_endpoint.Endpoint(
        name=_TEST_ENDPOINT_NAME,
        display_name="aaa",
        create_time=datetime.now() - timedelta(minutes=10),
    ),
]

_TEST_LIST_FILTER = 'display_name="abc"'
_TEST_LIST_ORDER_BY_CREATE_TIME = "create_time desc"
_TEST_LIST_ORDER_BY_DISPLAY_NAME = "display_name"

_TEST_LABELS = {"my_key": "my_value"}


@pytest.fixture
def get_endpoint_mock():
    with mock.patch.object(
        endpoint_service_client.EndpointServiceClient, "get_endpoint"
    ) as get_endpoint_mock:
        get_endpoint_mock.return_value = gca_endpoint.Endpoint(
            display_name=_TEST_DISPLAY_NAME,
            name=_TEST_ENDPOINT_NAME,
            encryption_spec=_TEST_ENCRYPTION_SPEC,
        )
        yield get_endpoint_mock


@pytest.fixture
def get_endpoint_alt_location_mock():
    with mock.patch.object(
        endpoint_service_client.EndpointServiceClient, "get_endpoint"
    ) as get_endpoint_mock:
        get_endpoint_mock.return_value = gca_endpoint.Endpoint(
            display_name=_TEST_DISPLAY_NAME,
            name=_TEST_ENDPOINT_NAME_ALT_LOCATION,
            encryption_spec=_TEST_ENCRYPTION_SPEC,
        )
        yield get_endpoint_mock


@pytest.fixture
def get_endpoint_with_models_mock():
    with mock.patch.object(
        endpoint_service_client.EndpointServiceClient, "get_endpoint"
    ) as get_endpoint_mock:
        get_endpoint_mock.return_value = gca_endpoint.Endpoint(
            display_name=_TEST_DISPLAY_NAME,
            name=_TEST_ENDPOINT_NAME,
            deployed_models=_TEST_DEPLOYED_MODELS,
        )
        yield get_endpoint_mock


@pytest.fixture
def get_model_mock():
    with mock.patch.object(
        model_service_client.ModelServiceClient, "get_model"
    ) as get_model_mock:
        get_model_mock.return_value = gca_model.Model(
            display_name=_TEST_DISPLAY_NAME, name=_TEST_MODEL_NAME,
        )
        yield get_model_mock


@pytest.fixture
def create_endpoint_mock():
    with mock.patch.object(
        endpoint_service_client.EndpointServiceClient, "create_endpoint"
    ) as create_endpoint_mock:
        create_endpoint_lro_mock = mock.Mock(ga_operation.Operation)
        create_endpoint_lro_mock.result.return_value = gca_endpoint.Endpoint(
            name=_TEST_ENDPOINT_NAME, display_name=_TEST_DISPLAY_NAME
        )
        create_endpoint_mock.return_value = create_endpoint_lro_mock
        yield create_endpoint_mock


@pytest.fixture
def deploy_model_mock():
    with mock.patch.object(
        endpoint_service_client.EndpointServiceClient, "deploy_model"
    ) as deploy_model_mock:
        deployed_model = gca_endpoint.DeployedModel(
            model=_TEST_MODEL_NAME, display_name=_TEST_DISPLAY_NAME,
        )
        deploy_model_lro_mock = mock.Mock(ga_operation.Operation)
        deploy_model_lro_mock.result.return_value = gca_endpoint_service.DeployModelResponse(
            deployed_model=deployed_model,
        )
        deploy_model_mock.return_value = deploy_model_lro_mock
        yield deploy_model_mock


@pytest.fixture
def deploy_model_with_explanations_mock():
    with mock.patch.object(
        endpoint_service_client_v1beta1.EndpointServiceClient, "deploy_model"
    ) as deploy_model_mock:
        deployed_model = gca_endpoint_v1beta1.DeployedModel(
            model=_TEST_MODEL_NAME, display_name=_TEST_DISPLAY_NAME,
        )
        deploy_model_lro_mock = mock.Mock(ga_operation.Operation)
        deploy_model_lro_mock.result.return_value = gca_endpoint_service_v1beta1.DeployModelResponse(
            deployed_model=deployed_model,
        )
        deploy_model_mock.return_value = deploy_model_lro_mock
        yield deploy_model_mock


@pytest.fixture
def undeploy_model_mock():
    with mock.patch.object(
        endpoint_service_client.EndpointServiceClient, "undeploy_model"
    ) as undeploy_model_mock:
        undeploy_model_lro_mock = mock.Mock(ga_operation.Operation)
        undeploy_model_lro_mock.result.return_value = (
            gca_endpoint_service.UndeployModelResponse()
        )
        undeploy_model_mock.return_value = undeploy_model_lro_mock
        yield undeploy_model_mock


@pytest.fixture
def delete_endpoint_mock():
    with mock.patch.object(
        endpoint_service_client.EndpointServiceClient, "delete_endpoint"
    ) as delete_endpoint_mock:
        delete_endpoint_lro_mock = mock.Mock(ga_operation.Operation)
        delete_endpoint_lro_mock.result.return_value = (
            gca_endpoint_service.DeleteEndpointRequest()
        )
        delete_endpoint_mock.return_value = delete_endpoint_lro_mock
        yield delete_endpoint_mock


@pytest.fixture
def sdk_private_undeploy_mock():
    """Mocks the high-level Endpoint._undeploy() SDK private method"""
    with mock.patch.object(aiplatform.Endpoint, "_undeploy") as sdk_undeploy_mock:
        sdk_undeploy_mock.return_value = None
        yield sdk_undeploy_mock


@pytest.fixture
def sdk_undeploy_all_mock():
    """Mocks the high-level Endpoint.undeploy_all() SDK method"""
    with mock.patch.object(
        aiplatform.Endpoint, "undeploy_all"
    ) as sdk_undeploy_all_mock:
        sdk_undeploy_all_mock.return_value = None
        yield sdk_undeploy_all_mock


@pytest.fixture
def list_endpoints_mock():
    with mock.patch.object(
        endpoint_service_client.EndpointServiceClient, "list_endpoints"
    ) as list_endpoints_mock:
        list_endpoints_mock.return_value = _TEST_ENDPOINT_LIST
        yield list_endpoints_mock


@pytest.fixture
def create_endpoint_client_mock():
    with mock.patch.object(
        initializer.global_config, "create_client", autospec=True,
    ) as create_endpoint_client_mock:
        endpoint_client_mock = mock.Mock(
            spec=endpoint_service_client.EndpointServiceClient
        )
        endpoint_client_mock.get_endpoint.return_value = _TEST_ENDPOINT_GAPIC
        create_endpoint_client_mock.return_value = endpoint_client_mock
        yield create_endpoint_client_mock


@pytest.fixture
def predict_client_predict_mock():
    with mock.patch.object(
        prediction_service_client.PredictionServiceClient, "predict"
    ) as predict_mock:
        predict_mock.return_value = gca_prediction_service.PredictResponse(
            deployed_model_id=_TEST_MODEL_ID
        )
        predict_mock.return_value.predictions.extend(_TEST_PREDICTION)
        yield predict_mock


@pytest.fixture
def predict_client_explain_mock():
    with mock.patch.object(
        prediction_service_client_v1beta1.PredictionServiceClient, "explain"
    ) as predict_mock:
        predict_mock.return_value = gca_prediction_service_v1beta1.ExplainResponse(
            deployed_model_id=_TEST_MODEL_ID,
        )
        predict_mock.return_value.predictions.extend(_TEST_PREDICTION)
        predict_mock.return_value.explanations.extend(_TEST_EXPLANATIONS)
        predict_mock.return_value.explanations[0].attributions.extend(
            _TEST_ATTRIBUTIONS
        )
        yield predict_mock


class TestEndpoint:
    def setup_method(self):
        reload(initializer)
        reload(aiplatform)
        aiplatform.init(project=_TEST_PROJECT, location=_TEST_LOCATION)

    def teardown_method(self):
        initializer.global_pool.shutdown(wait=True)

    def test_constructor(self, create_endpoint_client_mock):
        aiplatform.init(
            project=_TEST_PROJECT,
            location=_TEST_LOCATION,
            credentials=_TEST_CREDENTIALS,
        )
        models.Endpoint(_TEST_ENDPOINT_NAME)
        create_endpoint_client_mock.assert_has_calls(
            [
                mock.call(
                    client_class=utils.EndpointClientWithOverride,
                    credentials=initializer.global_config.credentials,
                    location_override=_TEST_LOCATION,
                    prediction_client=False,
                ),
                mock.call(
                    client_class=utils.PredictionClientWithOverride,
                    credentials=None,
                    location_override=_TEST_LOCATION,
                    prediction_client=True,
                ),
            ]
        )

    def test_constructor_with_endpoint_id(self, get_endpoint_mock):
        models.Endpoint(_TEST_ID)
        get_endpoint_mock.assert_called_with(name=_TEST_ENDPOINT_NAME)

    def test_constructor_with_endpoint_name(self, get_endpoint_mock):
        models.Endpoint(_TEST_ENDPOINT_NAME)
        get_endpoint_mock.assert_called_with(name=_TEST_ENDPOINT_NAME)

    def test_constructor_with_custom_project(self, get_endpoint_mock):
        models.Endpoint(endpoint_name=_TEST_ID, project=_TEST_PROJECT_2)
        test_endpoint_resource_name = endpoint_service_client.EndpointServiceClient.endpoint_path(
            _TEST_PROJECT_2, _TEST_LOCATION, _TEST_ID
        )
        get_endpoint_mock.assert_called_with(name=test_endpoint_resource_name)

    @pytest.mark.usefixtures("get_endpoint_mock")
    def test_constructor_with_conflicting_location(self):
        """Passing a full resource name with `_TEST_LOCATION` and providing `_TEST_LOCATION_2` as location"""

        with pytest.raises(RuntimeError) as err:
            models.Endpoint(
                endpoint_name=_TEST_ENDPOINT_NAME, location=_TEST_LOCATION_2
            )

        assert err.match(
            regexp=r"is provided, but different from the resource location"
        )

    def test_constructor_with_custom_location(self, get_endpoint_alt_location_mock):
        models.Endpoint(endpoint_name=_TEST_ID, location=_TEST_LOCATION_2)
        test_endpoint_resource_name = endpoint_service_client.EndpointServiceClient.endpoint_path(
            _TEST_PROJECT, _TEST_LOCATION_2, _TEST_ID
        )
        get_endpoint_alt_location_mock.assert_called_with(
            name=test_endpoint_resource_name
        )

    def test_constructor_with_custom_credentials(self, create_endpoint_client_mock):
        creds = auth_credentials.AnonymousCredentials()

        models.Endpoint(_TEST_ENDPOINT_NAME, credentials=creds)
        create_endpoint_client_mock.assert_has_calls(
            [
                mock.call(
                    client_class=utils.EndpointClientWithOverride,
                    credentials=creds,
                    location_override=_TEST_LOCATION,
                    prediction_client=False,
                ),
                mock.call(
                    client_class=utils.PredictionClientWithOverride,
                    credentials=creds,
                    location_override=_TEST_LOCATION,
                    prediction_client=True,
                ),
            ]
        )

    @pytest.mark.usefixtures("get_endpoint_mock")
    @pytest.mark.parametrize("sync", [True, False])
    def test_init_aiplatform_with_encryption_key_name_and_create_endpoint(
        self, create_endpoint_mock, sync
    ):
        aiplatform.init(
            project=_TEST_PROJECT,
            location=_TEST_LOCATION,
            encryption_spec_key_name=_TEST_ENCRYPTION_KEY_NAME,
        )
        my_endpoint = models.Endpoint.create(display_name=_TEST_DISPLAY_NAME, sync=sync)

        if not sync:
            my_endpoint.wait()

        expected_endpoint = gca_endpoint.Endpoint(
            display_name=_TEST_DISPLAY_NAME, encryption_spec=_TEST_ENCRYPTION_SPEC
        )
        create_endpoint_mock.assert_called_once_with(
            parent=_TEST_PARENT, endpoint=expected_endpoint, metadata=(),
        )

        expected_endpoint.name = _TEST_ENDPOINT_NAME
        assert my_endpoint._gca_resource == expected_endpoint

    @pytest.mark.usefixtures("get_endpoint_mock")
    @pytest.mark.parametrize("sync", [True, False])
    def test_create(self, create_endpoint_mock, sync):
        my_endpoint = models.Endpoint.create(
            display_name=_TEST_DISPLAY_NAME,
            encryption_spec_key_name=_TEST_ENCRYPTION_KEY_NAME,
            sync=sync,
        )

        if not sync:
            my_endpoint.wait()

        expected_endpoint = gca_endpoint.Endpoint(
            display_name=_TEST_DISPLAY_NAME, encryption_spec=_TEST_ENCRYPTION_SPEC
        )
        create_endpoint_mock.assert_called_once_with(
            parent=_TEST_PARENT, endpoint=expected_endpoint, metadata=(),
        )

        expected_endpoint.name = _TEST_ENDPOINT_NAME
        assert my_endpoint.gca_resource == expected_endpoint
        assert my_endpoint.network is None

    @pytest.mark.usefixtures("get_endpoint_mock")
    def test_accessing_properties_with_no_resource_raises(self,):

        my_endpoint = aiplatform.Endpoint(_TEST_ENDPOINT_NAME)

        my_endpoint._gca_resource = None

        with pytest.raises(RuntimeError) as e:
            my_endpoint.gca_resource
        e.match(regexp=r"Endpoint resource has not been created.")

        with pytest.raises(RuntimeError) as e:
            my_endpoint.network
        e.match(regexp=r"Endpoint resource has not been created.")

    @pytest.mark.usefixtures("get_endpoint_mock")
    @pytest.mark.parametrize("sync", [True, False])
    def test_create_with_description(self, create_endpoint_mock, sync):
        my_endpoint = models.Endpoint.create(
            display_name=_TEST_DISPLAY_NAME, description=_TEST_DESCRIPTION, sync=sync
        )
        if not sync:
            my_endpoint.wait()

        expected_endpoint = gca_endpoint.Endpoint(
            display_name=_TEST_DISPLAY_NAME, description=_TEST_DESCRIPTION,
        )
        create_endpoint_mock.assert_called_once_with(
            parent=_TEST_PARENT, endpoint=expected_endpoint, metadata=(),
        )

    @pytest.mark.usefixtures("get_endpoint_mock")
    @pytest.mark.parametrize("sync", [True, False])
    def test_create_with_labels(self, create_endpoint_mock, sync):
        my_endpoint = models.Endpoint.create(
            display_name=_TEST_DISPLAY_NAME, labels=_TEST_LABELS, sync=sync
        )
        if not sync:
            my_endpoint.wait()

        expected_endpoint = gca_endpoint.Endpoint(
            display_name=_TEST_DISPLAY_NAME, labels=_TEST_LABELS,
        )
        create_endpoint_mock.assert_called_once_with(