/*
 * Copyright 2004-present the original author or 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.
 */

package org.springframework.security.oauth2.jwt;

import java.util.Collections;
import java.util.List;
import java.util.function.Supplier;

import com.nimbusds.jose.jwk.JWK;
import com.nimbusds.jose.jwk.JWKMatcher;
import com.nimbusds.jose.jwk.JWKSelector;
import com.nimbusds.jose.jwk.KeyType;
import com.nimbusds.jose.jwk.KeyUse;
import okhttp3.mockwebserver.MockResponse;
import okhttp3.mockwebserver.MockWebServer;
import org.junit.jupiter.api.BeforeEach;
import org.junit.jupiter.api.Test;
import org.junit.jupiter.api.extension.ExtendWith;
import org.mockito.Mock;
import org.mockito.junit.jupiter.MockitoExtension;
import reactor.core.publisher.Mono;

import org.springframework.web.reactive.function.client.WebClientResponseException;

import static org.assertj.core.api.Assertions.assertThat;
import static org.assertj.core.api.Assertions.assertThatExceptionOfType;
import static org.mockito.ArgumentMatchers.any;
import static org.mockito.BDDMockito.given;
import static org.mockito.BDDMockito.willReturn;
import static org.mockito.BDDMockito.willThrow;

/**
 * @author Rob Winch
 * @since 5.1
 */
@ExtendWith(MockitoExtension.class)
public class ReactiveRemoteJWKSourceTests {

	@Mock
	private JWKMatcher matcher;

	private ReactiveRemoteJWKSource source;

	private JWKSelector selector;

	private MockWebServer server;

	@Mock
	private Supplier<String> mockStringSupplier;

	// @formatter:off
	private String keys = "{\n"
			+ "    \"keys\": [\n"
			+ "        {\n"
			+ "            \"alg\": \"RS256\", \n"
			+ "            \"e\": \"AQAB\", \n"
			+ "            \"kid\": \"1923397381d9574bb873202a90c32b7ceeaed027\", \n"
			+ "            \"kty\": \"RSA\", \n"
			+ "            \"n\": \"m4I5Dk5GnbzzUtqaljDVbpMONi1JLNJ8ZuXE8VvjCAVebDg5vTYhQ33jUwGgbn1wFmytUMgMmvK8A8Gpshl0sO2GBIZoh6_pwLrk657ZEtv-hx9fYKnzwyrfHqxtSswMAyr7XtKl8Ha1I03uFMSaYaaBTwVXCHByhzr4PVXfKAYJNbbcteUZfE8ODlBQkjQLI0IB78Nu8XIRrdzTF_5LCuM6rLUNtX6_KdzPpeX9KEtB7OBAfkdZEtBzGI-aYNLtIaL4qO6cVxBeVDLMoj9kVsRPylrwhEFQcGOjtJhwJwXFzTMZVhkiLFCHxZkkjoMrK5osSRlhduuGI9ot8XTUKQ\", \n"
			+ "            \"use\": \"sig\"\n"
			+ "        }, \n"
			+ "        {\n"
			+ "            \"alg\": \"RS256\", \n"
			+ "            \"e\": \"AQAB\", \n"
			+ "            \"kid\": \"7ddf54d3032d1f0d48c3618892ca74c1ac30ad77\", \n"
			+ "            \"kty\": \"RSA\", \n"
			+ "            \"n\": \"yLlYyux949b7qS-DdqTNjdZb4NtqiNH-Jt7DtRxmfW9XZLOQ6Q2NYgmPe9hyy5GHG7W3zsd6Q-rzq5eGRNEUx1767K1dS5PtkVWPiPG_M7rDqCu3HsLmKQKhRjHYaCWl5NuiMB5mXoPhSwrHd2yeGE7QHIV7_CiQFc1xQsXeiC-nTeJohJO3HI97w0GXE8pHspLYq9oG87f5IHxFr89abmwRug-D7QWQyW5b4doe4ZL-52J-8WHd52kGrGfu4QyV83oAad3I_9Q-yiWOXUr_0GIrzz4_-u5HgqYexnodFhZZSaKuRSg_b5qCnPhW8gBDLAHkmQzQMaWsN14L0pokbQ\", \n"
			+ "            \"use\": \"sig\"\n"
			+ "        }\n"
			+ "    ]\n"
			+ "}\n";
	// @formatter:on

	// @formatter:off
	private String keys2 = "{\n"
			+ "    \"keys\": [\n"
			+ "        {\n"
			+ "            \"alg\": \"RS256\", \n"
			+ "            \"e\": \"AQAB\", \n"
			+ "            \"kid\": \"rotated\", \n"
			+ "            \"kty\": \"RSA\", \n"
			+ "            \"n\": \"m4I5Dk5GnbzzUtqaljDVbpMONi1JLNJ8ZuXE8VvjCAVebDg5vTYhQ33jUwGgbn1wFmytUMgMmvK8A8Gpshl0sO2GBIZoh6_pwLrk657ZEtv-hx9fYKnzwyrfHqxtSswMAyr7XtKl8Ha1I03uFMSaYaaBTwVXCHByhzr4PVXfKAYJNbbcteUZfE8ODlBQkjQLI0IB78Nu8XIRrdzTF_5LCuM6rLUNtX6_KdzPpeX9KEtB7OBAfkdZEtBzGI-aYNLtIaL4qO6cVxBeVDLMoj9kVsRPylrwhEFQcGOjtJhwJwXFzTMZVhkiLFCHxZkkjoMrK5osSRlhduuGI9ot8XTUKQ\", \n"
			+ "            \"use\": \"sig\"\n"
			+ "        }\n"
			+ "    ]\n"
			+ "}\n";
	// @formatter:on

	@BeforeEach
	public void setup() {
		this.server = new MockWebServer();
		this.source = new ReactiveRemoteJWKSource(this.server.url("/").toString());
		this.server.enqueue(new MockResponse().setBody(this.keys));
		this.selector = new JWKSelector(this.matcher);
	}

	@Test
	public void getWhenMultipleRequestThenCached() {
		given(this.matcher.matches(any())).willReturn(true);
		this.source.get(this.selector).block();
		this.source.get(this.selector).block();
		assertThat(this.server.getRequestCount()).isEqualTo(1);
	}

	@Test
	public void getWhenMatchThenCreatesKeys() {
		given(this.matcher.matches(any())).willReturn(true);
		List<JWK> keys = this.source.get(this.selector).block();
		assertThat(keys).hasSize(2);
		JWK key1 = keys.get(0);
		assertThat(key1.getKeyID()).isEqualTo("1923397381d9574bb873202a90c32b7ceeaed027");
		assertThat(key1.getAlgorithm().getName()).isEqualTo("RS256");
		assertThat(key1.getKeyType()).isEqualTo(KeyType.RSA);
		assertThat(key1.getKeyUse()).isEqualTo(KeyUse.SIGNATURE);
		JWK key2 = keys.get(1);
		assertThat(key2.getKeyID()).isEqualTo("7ddf54d3032d1f0d48c3618892ca74c1ac30ad77");
		assertThat(key2.getAlgorithm().getName()).isEqualTo("RS256");
		assertThat(key2.getKeyType()).isEqualTo(KeyType.RSA);
		assertThat(key2.getKeyUse()).isEqualTo(KeyUse.SIGNATURE);
	}

	@Test
	public void getWhenNoMatchAndNoKeyIdThenEmpty() {
		given(this.matcher.matches(any())).willReturn(false);
		given(this.matcher.getKeyIDs()).willReturn(Collections.emptySet());
		assertThat(this.source.get(this.selector).block()).isEmpty();
	}

	@Test
	public void getWhenNoMatchAndKeyIdNotMatchThenRefreshAndFoundThenFound() {
		this.server.enqueue(new MockResponse().setBody(this.keys2));
		given(this.matcher.matches(any())).willReturn(false, false, true);
		given(this.matcher.getKeyIDs()).willReturn(Collections.singleton("rotated"));
		List<JWK> keys = this.source.get(this.selector).block();
		assertThat(keys).hasSize(1);
		assertThat(keys.get(0).getKeyID()).isEqualTo("rotated");
	}

	@Test
	public void getWhenNoMatchAndKeyIdNotMatchThenRefreshAndNotFoundThenEmpty() {
		this.server.enqueue(new MockResponse().setBody(this.keys2));
		given(this.matcher.matches(any())).willReturn(false, false, false);
		given(this.matcher.getKeyIDs()).willReturn(Collections.singleton("rotated"));
		List<JWK> keys = this.source.get(this.selector).block();
		assertThat(keys).isEmpty();
	}

	@Test
	public void getWhenNoMatchAndKeyIdMatchThenEmpty() {
		given(this.matcher.matches(any())).willReturn(false);
		given(this.matcher.getKeyIDs()).willReturn(Collections.singleton("7ddf54d3032d1f0d48c3618892ca74c1ac30ad77"));
		assertThat(this.source.get(this.selector).block()).isEmpty();
	}

	@Test
	public void getShouldRecoverAndReturnKeysAfterErrorCase() {
		given(this.matcher.matches(any())).willReturn(true);
		this.source = new ReactiveRemoteJWKSource(Mono.fromSupplier(this.mockStringSupplier));
		willThrow(WebClientResponseException.ServiceUnavailable.class).given(this.mockStringSupplier).get();
		// first case: id provider has error state
		assertThatExceptionOfType(WebClientResponseException.ServiceUnavailable.class)
			.isThrownBy(() -> this.source.get(this.selector).block());
		// second case: id provider is healthy again
		willReturn(this.server.url("/").toString()).given(this.mockStringSupplier).get();
		List<JWK> actual = this.source.get(this.selector).block();
		assertThat(actual).isNotEmpty();
	}

}
