# Copyright 2021 The KServe Authors.
#
# 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 json
import tempfile
from unittest import mock

import pytest
from kubernetes.client import V1ObjectMeta, V1ServiceAccount, V1ServiceAccountList, rest

from kserve import constants
from kserve.api.creds_utils import (
    check_sa_exists,
    create_secret,
    create_service_account,
    get_creds_name_from_config_map,
    patch_service_account,
    set_azure_credentials,
    set_gcs_credentials,
    set_s3_credentials,
    set_service_account,
)


@mock.patch("kubernetes.client.CoreV1Api.list_namespaced_service_account")
def test_check_sa_exists(mock_client):
    # Mock kubernetes client to return 2 accounts
    accounts = V1ServiceAccountList(
        items=[V1ServiceAccount(metadata=V1ObjectMeta(name=n)) for n in ["a", "b"]]
    )
    mock_client.return_value = accounts

    # then a, b should exist, c should not
    assert check_sa_exists("kubeflow", "a") is True
    assert check_sa_exists("kubeflow", "b") is True
    assert check_sa_exists("kubeflow", "c") is False


@mock.patch("kubernetes.client.CoreV1Api.create_namespaced_service_account")
def test_create_service_account(mock_client):
    sa_name = "test"
    namespace = "kserve-test"
    secret_name = "test_secret"
    create_service_account(secret_name, namespace, sa_name)
    mock_client.assert_called_once()

    mock_client.side_effect = rest.ApiException("foo")
    with pytest.raises(RuntimeError):
        sa_name = "test"
        namespace = "kserve-test"
        secret_name = "test_secret"
        create_service_account(secret_name, namespace, sa_name)


@mock.patch("kubernetes.client.CoreV1Api.patch_namespaced_service_account")
def test_patch_service_account(mock_client):
    sa_name = "test"
    namespace = "kserve-test"
    secret_name = "test_secret"
    patch_service_account(secret_name, namespace, sa_name)
    mock_client.assert_called_once()

    mock_client.side_effect = rest.ApiException("foo")
    with pytest.raises(RuntimeError):
        sa_name = "test"
        namespace = "kserve-test"
        secret_name = "test_secret"
        patch_service_account(secret_name, namespace, sa_name)


@mock.patch("kubernetes.client.CoreV1Api.create_namespaced_secret")
def test_create_secret(mock_create_secret):
    namespace = "test"
    secret_name = "test-secret"
    mock_create_secret.return_value = mock.Mock(**{"metadata.name": secret_name})
    assert create_secret(namespace) == secret_name

    with pytest.raises(RuntimeError):
        mock_create_secret.side_effect = rest.ApiException("foo")
        create_secret(namespace)


@mock.patch("kserve.api.creds_utils.create_service_account")
@mock.patch("kserve.api.creds_utils.patch_service_account")
@mock.patch("kserve.api.creds_utils.check_sa_exists")
def test_set_service_account(
    mock_check_sa_exists, mock_patch_service_account, mock_create_service_account
):
    namespace = "test"
    service_account = V1ServiceAccount()
    secret_name = "test-secret"
    mock_check_sa_exists.return_value = True
    set_service_account(namespace, service_account, secret_name)
    mock_patch_service_account.assert_called_once()

    mock_check_sa_exists.return_value = False
    set_service_account(namespace, service_account, secret_name)
    mock_create_service_account.assert_called_once()


@mock.patch("kubernetes.client.CoreV1Api.read_namespaced_config_map")
def test_get_creds_name_from_config_map(mock_read_config_map):
    mock_read_config_map.return_value = mock.Mock(
        **{
            "data": {
                "credentials": """{
        "gcs": {"gcsCredentialFileName": "gcs_cred.json"},
        "s3": {"s3AccessKeyIDName": "s3_access_key.json",
               "s3SecretAccessKeyName": "s3_secret.json"}}"""
            }
        }
    )
    test_cases = {
        "gcsCredentialFileName": "gcs_cred.json",
        "s3AccessKeyIDName": "s3_access_key.json",
        "s3SecretAccessKeyName": "s3_secret.json",
    }
    for cred, result in test_cases.items():
        assert get_creds_name_from_config_map(cred) == result

    with pytest.raises(RuntimeError):
        get_creds_name_from_config_map("invalidCred")

    mock_read_config_map.side_effect = rest.ApiException("foo")
    assert get_creds_name_from_config_map("gcsCredentialFileName") is None


@mock.patch("kserve.api.creds_utils.set_service_account")
@mock.patch("kserve.api.creds_utils.create_secret")
@mock.patch("kserve.api.creds_utils.get_creds_name_from_config_map")
def test_set_gcs_credentials(
    mock_get_creds_name, mock_create_secret, mock_set_service_account
):
    namespace = "test"
    service_account = V1ServiceAccount()
    temp_cred_file = tempfile.NamedTemporaryFile(suffix=".json")
    cred_file_name = temp_cred_file.name
    mock_get_creds_name.return_value = cred_file_name
    mock_create_secret.return_value = "test-secret"
    set_gcs_credentials(namespace, cred_file_name, service_account)
    mock_get_creds_name.assert_called()
    mock_create_secret.assert_called()
    mock_set_service_account.assert_called()

    mock_get_creds_name.return_value = None
    set_gcs_credentials(namespace, cred_file_name, service_account)
    mock_get_creds_name.assert_called()
    mock_create_secret.assert_called()
    mock_set_service_account.assert_called()


@mock.patch("kserve.api.creds_utils.set_service_account")
@mock.patch("kserve.api.creds_utils.create_secret")
@mock.patch("kserve.api.creds_utils.get_creds_name_from_config_map")
def test_set_s3_credentials(
    mock_get_creds_name, mock_create_secret, mock_set_service_account
):
    namespace = "test"
    endpoint = "https://s3.aws.com"
    region = "ap-south-1"
    use_https = True
    verfify_ssl = True
    cabundle = "/user/test/cert.pem"
    data = {
        constants.S3_ACCESS_KEY_ID_DEFAULT_NAME: "XXXXXXXXXXXX",
        constants.S3_SECRET_ACCESS_KEY_DEFAULT_NAME: "XXXXXXXXXXXX",
    }
    annotations = {
        constants.KSERVE_GROUP + "/s3-endpoint": endpoint,
        constants.KSERVE_GROUP + "/s3-region": region,
        constants.KSERVE_GROUP + "/s3-usehttps": use_https,
        constants.KSERVE_GROUP + "/s3-verifyssl": verfify_ssl,
        constants.KSERVE_GROUP + "/s3-cabundle": cabundle,
    }
    creds_str = b"""
    [default]
    aws_access_key_id = XXXXXXXXXXXX
    aws_secret_access_key = XXXXXXXXXXXX
    """

    with tempfile.NamedTemporaryFile() as creds_file:
        creds_file.write(creds_str)
        creds_file.seek(0)
        mock_get_creds_name.return_value = None
        mock_create_secret.return_value = "test-secret"
        set_s3_credentials(
            namespace,
            creds_file.name,
            V1ServiceAccount(),
            s3_endpoint=endpoint,
            s3_region=region,
            s3_use_https=use_https,
            s3_verify_ssl=verfify_ssl,
            s3_cabundle=cabundle,
        )
    mock_create_secret.assert_called_with(
        namespace=namespace, annotations=annotations, data=data
    )
    mock_get_creds_name.asset_called()
    mock_set_service_account.assert_called()


@mock.patch("kserve.api.creds_utils.set_service_account")
@mock.patch("kserve.api.creds_utils.create_secret")
def test_set_azure_credentials(mock_create_secret, mock_set_service_account):
    namespace = "test"
    creds = {
        "clientId": "XXXXXXXXXXX",
        "clientSecret": "XXXXXXXXXXX",
        "subscriptionId": "XXXXXXXXXXX",
        "tenantId": "XXXXXXXXXXX",
    }
    data = {
        "AZURE_CLIENT_ID": creds["clientId"],
        "AZURE_CLIENT_SECRET": creds["clientSecret"],
        "AZURE_SUBSCRIPTION_ID": creds["subscriptionId"],
        "AZURE_TENANT_ID": creds["tenantId"],
    }
    with tempfile.NamedTemporaryFile(suffix=".json") as creds_file:
        creds_file.write(json.dumps(creds).encode("utf-8"))
        creds_file.seek(0)
        mock_create_secret.return_value = "test-secret"
        set_azure_credentials(namespace, creds_file.name, V1ServiceAccount())
    mock_create_secret.assert_called_with(namespace=namespace, data=data)
    mock_set_service_account.assert_called()
