# SPDX-License-Identifier: Apache-2.0
# Copyright 2018-2020 The Feast 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
#
#     https://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
from http import HTTPStatus
from unittest.mock import call, patch

from pytest import fixture, raises

from feast.config import Config
from feast.grpc.auth import (
    GoogleOpenIDAuthMetadataPlugin,
    OAuthMetadataPlugin,
    get_auth_metadata_plugin,
)

AUDIENCE = "https://testaudience.io/"

AUTH_URL = "https://test.auth.com/v2/token"

HEADERS = {"content-type": "application/json"}

DATA = json.dumps(
    {
        "grant_type": "client_credentials",
        "client_id": "fakeID",
        "client_secret": "fakeSecret",
        "audience": AUDIENCE,
    }
)


class MockResponse:
    def __init__(self, json_data, status_code):
        self.json_data = json_data
        self.status_code = status_code

    def json(self):
        return self.json_data


class GoogleMockResponse:
    def __init__(self, stdout):
        self.stdout = stdout


class GoogleDefaultResponse:
    def __init__(self, id_token):
        self.id_token = id_token

    def refresh(self, request):
        pass


class GoogleDefaultErrorResponse:
    def __init__(self, id_token):
        self.wrong_attribute = id_token

    def refresh(self, request):
        pass


@fixture
def config_oauth():
    config_dict = {
        "core_url": "localhost:50051",
        "enable_auth": True,
        "auth_provider": "oauth",
        "oauth_grant_type": "client_credentials",
        "oauth_client_id": "fakeID",
        "oauth_client_secret": "fakeSecret",
        "oauth_audience": AUDIENCE,
        "oauth_token_request_url": AUTH_URL,
    }
    return Config(config_dict)


@fixture
def config_google():
    config_dict = {
        "core_url": "localhost:50051",
        "enable_auth": True,
        "auth_provider": "google",
    }
    return Config(config_dict)


@fixture
def config_with_missing_variable():
    config_dict = {
        "core_url": "localhost:50051",
        "enable_auth": True,
        "auth_provider": "oauth",
        "oauth_grant_type": "client_credentials",
        "oauth_client_id": "fakeID",
        "oauth_client_secret": "fakeSecret",
        "oauth_token_request_url": AUTH_URL,
    }
    return Config(config_dict)


@patch(
    "requests.post",
    return_value=MockResponse({"access_token": "mock_token"}, HTTPStatus.OK),
)
def test_get_auth_metadata_plugin_oauth_should_pass(post, config_oauth):
    auth_metadata_plugin = get_auth_metadata_plugin(config_oauth)
    assert isinstance(auth_metadata_plugin, OAuthMetadataPlugin)
    assert post.call_count == 1
    assert post.call_args == call(AUTH_URL, headers=HEADERS, data=DATA)
    assert auth_metadata_plugin.get_signed_meta() == (
        ("authorization", "Bearer mock_token"),
    )


@patch(
    "requests.post",
    return_value=MockResponse({"access_token": "mock_token"}, HTTPStatus.UNAUTHORIZED),
)
def test_get_auth_metadata_plugin_oauth_should_raise_when_response_is_not_200(
    post, config_oauth
):
    with raises(RuntimeError):
        get_auth_metadata_plugin(config_oauth)
        assert post.call_count == 1
        assert post.call_args == call(AUTH_URL, headers=HEADERS, data=DATA)


def test_get_auth_metadata_plugin_oauth_should_raise_when_config_is_incorrect(
    config_with_missing_variable,
):
    with raises(RuntimeError):
        get_auth_metadata_plugin(config_with_missing_variable)


@patch("google.auth.jwt.decode", return_value=GoogleMockResponse("jwt_token"))
@patch(
    "subprocess.run", return_value=GoogleMockResponse("std_output".encode("utf-8")),
)
def test_get_auth_metadata_plugin_google_should_pass_with_token_from_gcloud_sdk(
    subprocess, jwt, config_google
):
    auth_metadata_plugin = get_auth_metadata_plugin(config_google)
    assert isinstance(auth_metadata_plugin, GoogleOpenIDAuthMetadataPlugin)
    assert auth_metadata_plugin.get_signed_meta() == (
        ("authorization", "Bearer std_output"),
    )


@patch(
    "google.auth.default",
    return_value=[
        GoogleDefaultResponse("fake_token"),
        GoogleDefaultResponse("project_id"),
    ],
)
@patch(
    "subprocess.run", return_value=GoogleMockResponse("std_output".encode("utf-8")),
)
def test_get_auth_metadata_plugin_google_should_pass_with_token_from_google_auth_lib(
    subprocess, default, config_google
):
    auth_metadata_plugin = get_auth_metadata_plugin(config_google)
    assert isinstance(auth_metadata_plugin, GoogleOpenIDAuthMetadataPlugin)
    assert auth_metadata_plugin.get_signed_meta() == (
        ("authorization", "Bearer fake_token"),
    )


@patch(
    "google.auth.default",
    return_value=[
        GoogleDefaultErrorResponse("fake_token"),
        GoogleDefaultErrorResponse("project_id"),
    ],
)
@patch(
    "subprocess.run", return_value=GoogleMockResponse("std_output".encode("utf-8")),
)
def test_get_auth_metadata_plugin_google_should_raise_when_token_validation_fails(
    subprocess, default, config_google
):
    with raises(RuntimeError):
        get_auth_metadata_plugin(config_google)
