/*
 * Copyright 2002-2017 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
 *
 *      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.
 */

package org.springframework.web.reactive.function.client;

import java.nio.charset.CharsetEncoder;
import java.nio.charset.StandardCharsets;
import java.util.Base64;
import java.util.Map;
import java.util.Optional;
import java.util.function.Consumer;
import java.util.function.Function;
import java.util.function.Predicate;

import reactor.core.publisher.Mono;

import org.springframework.http.HttpHeaders;
import org.springframework.http.HttpStatus;
import org.springframework.util.Assert;

/**
 * Implementations of {@link ExchangeFilterFunction} that provide various useful request filter
 * operations, such as basic authentication, error handling, etc.
 *
 * @author Rob Winch
 * @author Arjen Poutsma
 * @since 5.0
 */
public abstract class ExchangeFilterFunctions {

	/**
	 * Name of the {@link ClientRequest} attribute that contains the
	 * {@link Credentials}, as used by {@link #basicAuthentication()}.
	 */
	public static final String BASIC_AUTHENTICATION_CREDENTIALS_ATTRIBUTE =
			ExchangeFilterFunctions.class.getName() + ".basicAuthenticationCredentials";


	/**
	 * Return a filter that adds an Authorization header for HTTP Basic Authentication, based on
	 * the given username and password.
	 * <p>Note that Basic Authentication only supports characters in the
	 * {@link StandardCharsets#ISO_8859_1 ISO-8859-1} character set.
	 * @param username the username to use
	 * @param password the password to use
	 * @return the {@link ExchangeFilterFunction} that adds the Authorization header
	 * @throws IllegalArgumentException if either {@code username} or {@code password} contain
	 * characters that cannot be encoded to ISO-8859-1
	 */
	public static ExchangeFilterFunction basicAuthentication(String username, String password) {
		Assert.notNull(username, "'username' must not be null");
		Assert.notNull(password, "'password' must not be null");

		checkIllegalCharacters(username, password);
		return basicAuthenticationInternal(r -> Optional.of(new Credentials(username, password)));
	}

	/**
	 * Return a filter that adds an Authorization header for HTTP Basic Authentication, based on
	 * the {@link Credentials} provided in the
	 * {@linkplain ClientRequest#attributes() request attributes}. If the attribute is not found,
	 * no authorization header is added.
	 * <p>Note that Basic Authentication only supports characters in the
	 * {@link StandardCharsets#ISO_8859_1 ISO-8859-1} character set.
	 * @return the {@link ExchangeFilterFunction} that adds the Authorization header
	 * @see #BASIC_AUTHENTICATION_CREDENTIALS_ATTRIBUTE
	 * @see Credentials#basicAuthenticationCredentials(String, String)
	 */
	public static ExchangeFilterFunction basicAuthentication() {
		return basicAuthenticationInternal(
				request -> request.attribute(BASIC_AUTHENTICATION_CREDENTIALS_ATTRIBUTE).map(o -> (Credentials)o));
	}

	private static ExchangeFilterFunction basicAuthenticationInternal(
			Function<ClientRequest, Optional<Credentials>> credentialsFunction) {

		return ExchangeFilterFunction.ofRequestProcessor(
				clientRequest -> credentialsFunction.apply(clientRequest).map(
						credentials -> {
							ClientRequest authorizedRequest = ClientRequest.from(clientRequest)
									.headers(headers -> headers.set(HttpHeaders.AUTHORIZATION,
                                            authorization(credentials)))
									.build();
							return Mono.just(authorizedRequest);
						})
						.orElse(Mono.just(clientRequest)));
	}

	private static String authorization(Credentials credentials) {
		String credentialsString = credentials.username + ":" + credentials.password;
		byte[] credentialBytes = credentialsString.getBytes(StandardCharsets.ISO_8859_1);
		byte[] encodedBytes = Base64.getEncoder().encode(credentialBytes);
		String encodedCredentials = new String(encodedBytes, StandardCharsets.ISO_8859_1);
		return "Basic " + encodedCredentials;
	}

	/*
	 * Basic authentication only supports ISO 8859-1, see
	 * https://stackoverflow.com/questions/702629/utf-8-characters-mangled-in-http-basic-auth-username#703341
	 */
	private static void checkIllegalCharacters(String username, String password) {
		CharsetEncoder encoder = StandardCharsets.ISO_8859_1.newEncoder();
		if (!encoder.canEncode(username) || !encoder.canEncode(password)) {
			throw new IllegalArgumentException(
					"Username or password contains characters that cannot be encoded to ISO-8859-1");
		}

	}



	/**
	 * Return a filter that returns a given {@link Throwable} as response if the given
	 * {@link HttpStatus} predicate matches.
	 * @param statusPredicate the predicate that should match the
	 * {@linkplain ClientResponse#statusCode() response status}
	 * @param exceptionFunction the function that returns the exception
	 * @return the {@link ExchangeFilterFunction} that returns the given exception if the predicate
	 * matches
	 */
	public static ExchangeFilterFunction statusError(Predicate<HttpStatus> statusPredicate,
			Function<ClientResponse, ? extends Throwable> exceptionFunction) {

		Assert.notNull(statusPredicate, "'statusPredicate' must not be null");
		Assert.notNull(exceptionFunction, "'exceptionFunction' must not be null");

		return ExchangeFilterFunction.ofResponseProcessor(
				clientResponse -> {
					if (statusPredicate.test(clientResponse.statusCode())) {
						return Mono.error(exceptionFunction.apply(clientResponse));
					}
					else {
						return Mono.just(clientResponse);
					}
				}
		);
	}


	/**
	 * Represents a combination of username and password, as used by {@link #basicAuthentication()}.
	 * @see #basicAuthenticationCredentials(String, String)
	 */
	public static final class Credentials {

		private final String username;

		private final String password;

		/**
		 * Create a new {@code Credentials} instance with the given username and password.
		 * @param username the username
		 * @param password the password
		 */
		public Credentials(String username, String password) {
			Assert.notNull(username, "'username' must not be null");
			Assert.notNull(password, "'password' must not be null");

			this.username = username;
			this.password = password;
		}

		/**
		 * Return a consumer that stores the given username and password in the
		 * {@linkplain ClientRequest.Builder#attributes(java.util.function.Consumer) request
		 * attributes} as a {@code Credentials} object.
		 * @param username the username
		 * @param password the password
		 * @return a consumer that adds the given credentials to the attribute map
		 * @see ClientRequest.Builder#attributes(java.util.function.Consumer)
		 * @see #BASIC_AUTHENTICATION_CREDENTIALS_ATTRIBUTE
		 */
		public static Consumer<Map<String, Object>> basicAuthenticationCredentials(String username, String password) {
			Credentials credentials = new Credentials(username, password);
			checkIllegalCharacters(username, password);

			return attributes -> attributes.put(BASIC_AUTHENTICATION_CREDENTIALS_ATTRIBUTE, credentials);
		}

		@Override
		public boolean equals(Object o) {
			if (this == o) {
				return true;
			}
			if (o instanceof Credentials) {
				Credentials other = (Credentials) o;
				return this.username.equals(other.username) &&
						this.password.equals(other.password);

			}
			return false;
		}

		@Override
		public int hashCode() {
			return 31 * this.username.hashCode() + this.password.hashCode();
		}

	}

}
