package cz.muni.ics.oauth2.token.exchange;

import com.fasterxml.jackson.core.JsonProcessingException;
import com.fasterxml.jackson.databind.JsonNode;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.fasterxml.jackson.databind.annotation.JsonSerialize;
import com.fasterxml.jackson.databind.node.JsonNodeFactory;
import com.fasterxml.jackson.databind.node.ObjectNode;
import cz.muni.ics.oauth2.model.OAuth2AccessTokenEntity;
import cz.muni.ics.oauth2.service.OAuth2TokenEntityService;
import cz.muni.ics.oauth2.token.TokenExchangeGranter;
import cz.muni.ics.oidc.server.configurations.PerunOidcConfig;
import lombok.Getter;
import lombok.ToString;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.security.oauth2.common.OAuth2AccessToken;
import org.springframework.security.oauth2.common.OAuth2AccessTokenJackson2Serializer;
import org.springframework.security.oauth2.common.OAuth2RefreshToken;
import org.springframework.security.oauth2.common.exceptions.InvalidRequestException;
import org.springframework.security.oauth2.common.exceptions.OAuth2Exception;
import org.springframework.security.oauth2.provider.ClientDetails;
import org.springframework.security.oauth2.provider.TokenRequest;
import org.springframework.stereotype.Component;
import org.springframework.util.StringUtils;

import java.net.Authenticator;
import java.net.PasswordAuthentication;
import java.net.URI;
import java.net.http.HttpClient;
import java.net.http.HttpRequest;
import java.net.http.HttpResponse;
import java.util.Date;
import java.util.HashSet;
import java.util.Map;
import java.util.Set;
import java.util.stream.Collectors;

import static cz.muni.ics.oauth2.token.TokenExchangeGranter.GRANT_TYPE;
import static cz.muni.ics.oauth2.token.TokenExchangeGranter.ISSUED_TOKEN_TYPE;
import static cz.muni.ics.oauth2.token.TokenExchangeGranter.TOKEN_TYPE_ACCESS_TOKEN;
import static org.springframework.http.HttpHeaders.CONTENT_TYPE;
import static org.springframework.http.MediaType.APPLICATION_JSON_VALUE;

@Component
@Slf4j
public class KerberosTicketTokenExchangeGranter extends BaseTokenExchangeGranter {

	private static final String TYPE = "type";

	private static final String TICKET = "ticket";

	private static final String EXPIRATION = "expiration";

	public static final String TOKEN_TYPE_KRB_TICKET = "urn:cesnet:proxyidp:oauth:token-type:kerberos_ticket";

	@Getter(onMethod_ = @Override)
	private final Set<String> supportedSubjectTokenTypes = Set.of(TOKEN_TYPE_ACCESS_TOKEN);

	@Getter(onMethod_ = @Override)
	private final Set<String> supportedRequestedTokenTypes = Set.of(TOKEN_TYPE_KRB_TICKET);

	@Getter(onMethod_ = @Override)
	private final Set<String> supportedActorTokenTypes = Set.of(TOKEN_TYPE_ACCESS_TOKEN);

	private final ObjectMapper objectMapper = new ObjectMapper();

	@Getter(onMethod_ = @Override)
	private final boolean supported;

	private final String exchangeApiUrl;

	private final String exchangeApiUsername;

	private final String exchangeApiPassword;

	private final Set<String> exchangeRequiredScopes;

	private final HttpClient httpClient;

	@Autowired
	protected KerberosTicketTokenExchangeGranter(PerunOidcConfig oidcConfig,
												 OAuth2TokenEntityService tokenServices,
												 TokenExchangeGranter tokenExchangeGranter)
	{
		super(tokenServices, tokenExchangeGranter);
		exchangeApiUrl = oidcConfig.getKrbTokenExchangeApiUrl();
		exchangeApiUsername = oidcConfig.getKrbTokenExchangeApiUsername();
		exchangeApiPassword = oidcConfig.getKrbTokenExchangeApiPassword();
		if (oidcConfig.getKrbTokenExchangeRequiredScopes() != null) {
			exchangeRequiredScopes = oidcConfig.getKrbTokenExchangeRequiredScopes().stream().filter(StringUtils::hasText).collect(Collectors.toSet());
		} else {
			exchangeRequiredScopes = new HashSet<>();
		}
		if (StringUtils.hasText(oidcConfig.getKrbTokenExchangeApiUrl())) {
			log.debug("Enabling KERBEROS token exchange support");
			HttpClient.Builder builder = HttpClient.newBuilder();
			if (StringUtils.hasText(exchangeApiUsername) && StringUtils.hasText(exchangeApiPassword)) {
				builder = builder.authenticator(new Authenticator() {
					@Override
					protected PasswordAuthentication getPasswordAuthentication() {
						return new PasswordAuthentication(exchangeApiUsername, exchangeApiPassword.toCharArray());
					}
				});
			}
			httpClient = builder.build();
			supported = true;
		} else {
			httpClient = null;
			supported = false;
		}
	}

	@Override
	protected boolean validateCallingClient(ClientDetails client) {
		return client.getAuthorizedGrantTypes().contains(GRANT_TYPE);
	}

	@Override
	protected boolean validateIncomingAccessToken(OAuth2AccessTokenEntity incomingToken) {
		if (!super.validateIncomingAccessToken(incomingToken)) {
			return false;
		};
		if (exchangeRequiredScopes != null && exchangeRequiredScopes.size() > 0) {
			Set<String> tokenScopes = incomingToken.getScope();
			for (String requiredScope: exchangeRequiredScopes) {
				if (!tokenScopes.contains(requiredScope)) {
					log.warn("Missing required scope '{}' in token to perform token exchange " +
							"(token scopes: {}, required scopes: {})", requiredScope, tokenScopes, exchangeRequiredScopes);
					throw new InvalidRequestException("Missing required scopes in token to perform the exchange - "
							+ exchangeRequiredScopes
					);
				}
			}
		}
		return true;
	}

	@Override
	protected OAuth2AccessToken getAccessToken(ClientDetails clientDetails,
											   OAuth2AccessTokenEntity accessToken,
											   String requestedTokenType,
											   TokenRequest tokenRequest)
	{
		try {
			ObjectNode body = JsonNodeFactory.instance.objectNode();
			body.put("token", accessToken.getValue());
			HttpRequest request = HttpRequest.newBuilder()
					.POST(HttpRequest.BodyPublishers.ofString(body.toPrettyString()))
					.header(CONTENT_TYPE, APPLICATION_JSON_VALUE)
					.uri(new URI(this.exchangeApiUrl))
					.build();
			log.debug("Contacting KRB exchanger endpoint '{}', caller client_id: {}, token subject: {}",
					this.exchangeApiUrl, clientDetails.getClientId(), accessToken.getJwtValue().getJWTClaimsSet().getSubject());
			HttpResponse<String> response = httpClient.send(request, HttpResponse.BodyHandlers.ofString());
			if (response.statusCode() == 200) {
				JsonNode responseJson;
				try {
					responseJson = objectMapper.readTree(response.body());
				} catch (JsonProcessingException e) {
					log.error("Could not parse JSON response from exchange", e);
					throw new Exception(e);
				}
				if (responseJson != null && responseJson.isObject()) {
					if (!responseJson.hasNonNull(TICKET)) {
						log.warn("Invalid response received - missing field '{}'", TICKET);
					} else if (!responseJson.hasNonNull(EXPIRATION)) {
						log.warn("Invalid response received - missing field '{}'", EXPIRATION);
					} else {
						log.debug("Transforming response from exchanger API to token endpoint response");
						KerberosTranslatedToken translatedToken = new KerberosTranslatedToken(responseJson);
						log.debug("Returning translated token '{}'", translatedToken);
						return translatedToken;
					}
				}
			} else {
				log.warn("Contacting KRB Token exchanger returned non-200 response code ({})", response.statusCode());
				log.debug("Response body: {}", response.body());
			}
		} catch (Exception ex) {
			log.warn("Caught exception when exchanging token for KRB Ticket", ex);
		}
		throw OAuth2Exception.create(OAuth2Exception.ERROR, "failed to execute token-exchange");
	}

	@Override
	public String toString() {
		return "KerberosTicketTokenExchangeGranter{" +
				"supportedSubjectTokenTypes=" + supportedSubjectTokenTypes +
				", supportedRequestedTokenTypes=" + supportedRequestedTokenTypes +
				", supportedActorTokenTypes=" + supportedActorTokenTypes +
				", supported=" + supported +
				", exchangeApiUrl='" + exchangeApiUrl + '\'' +
				", exchangeApiUsername='" + exchangeApiUsername + '\'' +
				", exchangeApiPassword='__PROTECTED__'" +
				'}';
	}

	@ToString
	@Getter(onMethod_ = @Override)
	@JsonSerialize(using = OAuth2AccessTokenJackson2Serializer.class)
	private static final class KerberosTranslatedToken implements OAuth2AccessToken {

		@ToString.Exclude
		private final String value;

		private final Date expiration;

		private final String tokenType;

		private final int expiresIn;

		private final boolean expired;

		KerberosTranslatedToken(JsonNode json) {
			if (json == null || json.isNull() || json.isEmpty() || !json.isObject()) {
				throw new IllegalArgumentException("Invalid JSON provided");
			} else if (!json.hasNonNull(TICKET)) {
				throw new IllegalArgumentException("Missing non-null field TYPE in JSON");
			} else if (!json.hasNonNull(EXPIRATION)) {
				throw new IllegalArgumentException("Missing non-null field EXPIRATION in JSON");
			}
			this.value = json.get(TICKET).asText();
			this.expiration = new Date(json.get(EXPIRATION).longValue() * 1000);
			this.tokenType = json.hasNonNull(TYPE) ? json.get(TYPE).textValue() : "Kerberos";
			if (json.hasNonNull(EXPIRES_IN)) {
				this.expiresIn = Math.toIntExact(json.get(EXPIRES_IN).longValue());
			} else {
				long secondsDiff = (expiration.getTime() - new Date().getTime()) / 1000;
				this.expiresIn = Math.toIntExact(Math.max(secondsDiff, 0));
			}
			this.expired = 0 == expiresIn;
		}

		@Override
		public Map<String, Object> getAdditionalInformation() {
			return Map.of(ISSUED_TOKEN_TYPE, TOKEN_TYPE_KRB_TICKET);
		}

		@Override
		public Set<String> getScope() {
			return null;
		}

		@Override
		public OAuth2RefreshToken getRefreshToken() {
			return null;
		}
	}

}
