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

import com.google.common.base.Joiner;
import com.google.common.base.Strings;
import com.google.common.collect.Lists;
import com.nimbusds.jose.JOSEObjectType;
import com.nimbusds.jose.JWSAlgorithm;
import com.nimbusds.jose.JWSHeader;
import com.nimbusds.jwt.JWTClaimsSet;
import com.nimbusds.jwt.SignedJWT;
import cz.muni.ics.jwt.signer.service.JWTSigningAndValidationService;
import cz.muni.ics.oauth2.exception.InvalidTargetException;
import cz.muni.ics.oauth2.model.AuthenticationHolderEntity;
import cz.muni.ics.oauth2.model.ClientDetailsEntity;
import cz.muni.ics.oauth2.model.OAuth2AccessTokenEntity;
import cz.muni.ics.oauth2.model.OAuth2RefreshTokenEntity;
import cz.muni.ics.oauth2.repository.AuthenticationHolderRepository;
import cz.muni.ics.oauth2.service.ClientDetailsEntityService;
import cz.muni.ics.oauth2.service.OAuth2TokenEntityService;
import cz.muni.ics.oauth2.token.TokenExchangeGranter;
import cz.muni.ics.oidc.server.configurations.PerunOidcConfig;
import cz.muni.ics.openid.connect.model.ApprovedSite;
import cz.muni.ics.openid.connect.request.ConnectRequestParameters;
import cz.muni.ics.openid.connect.service.ApprovedSiteService;
import lombok.Getter;
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.exceptions.InvalidClientException;
import org.springframework.security.oauth2.common.exceptions.InvalidScopeException;
import org.springframework.security.oauth2.common.exceptions.InvalidTokenException;
import org.springframework.security.oauth2.provider.ClientDetails;
import org.springframework.security.oauth2.provider.OAuth2Request;
import org.springframework.security.oauth2.provider.TokenRequest;
import org.springframework.stereotype.Component;
import org.springframework.util.StringUtils;

import java.text.ParseException;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Date;
import java.util.HashMap;
import java.util.HashSet;
import java.util.Map;
import java.util.Set;
import java.util.UUID;

import static cz.muni.ics.oauth2.service.IntrospectionResultAssembler.ACR;
import static cz.muni.ics.oauth2.service.IntrospectionResultAssembler.AMR;
import static cz.muni.ics.oauth2.service.IntrospectionResultAssembler.AUTH_TIME;
import static cz.muni.ics.oauth2.service.IntrospectionResultAssembler.CLIENT_ID;
import static cz.muni.ics.oauth2.service.IntrospectionResultAssembler.SCOPE;
import static cz.muni.ics.oauth2.service.IntrospectionResultAssembler.SCOPE_SEPARATOR;
import static cz.muni.ics.oauth2.service.SystemScopeService.OFFLINE_ACCESS;
import static cz.muni.ics.oauth2.token.TokenExchangeGranter.ISSUED_TOKEN_TYPE;
import static cz.muni.ics.oauth2.token.TokenExchangeGranter.TOKEN_TYPE_ACCESS_TOKEN;

@Component
@Slf4j
public class OAuthTokenExchangeGranter extends BaseTokenExchangeGranter {

	@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_ACCESS_TOKEN);

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

	private final ApprovedSiteService approvedSiteService;

	private final AuthenticationHolderRepository authenticationHolderRepository;

	private final JWTSigningAndValidationService jwtService;

	private final PerunOidcConfig config;

	private final ClientDetailsEntityService clientDetailsEntityService;

	@Autowired
	protected OAuthTokenExchangeGranter(OAuth2TokenEntityService tokenServices,
										TokenExchangeGranter tokenExchangeGranter,
										ApprovedSiteService approvedSiteService,
										AuthenticationHolderRepository authenticationHolderRepository,
										JWTSigningAndValidationService jwtService,
										PerunOidcConfig config,
										ClientDetailsEntityService clientDetailsEntityService)
	{
		super(tokenServices, tokenExchangeGranter);
		this.approvedSiteService = approvedSiteService;
		this.authenticationHolderRepository = authenticationHolderRepository;
		this.jwtService = jwtService;
		this.config = config;
		this.clientDetailsEntityService = clientDetailsEntityService;
	}

	@Override
	protected boolean validateCallingClient(ClientDetails client) {
		return true;
	}

	@Override
	protected OAuth2AccessToken getAccessToken(ClientDetails clientDetails,
											   OAuth2AccessTokenEntity subjectToken,
											   String requestedTokenType,
											   TokenRequest tokenRequest)
	{
		if (!(clientDetails instanceof ClientDetailsEntity)) {
			throw new InvalidClientException("Unrecognized client");
		}
		ClientDetailsEntity client = (ClientDetailsEntity) clientDetails;
		OAuth2TokenEntityService tokenServices = super.getTokenServices();

		OAuth2AccessTokenEntity token = new OAuth2AccessTokenEntity();

		token.setClient(client);
		Set<String> scopes = new HashSet<>();
		if (tokenRequest.getScope() != null && !tokenRequest.getScope().isEmpty()) {
			scopes.addAll(tokenRequest.getScope());
		} else {
			scopes.addAll(subjectToken.getScope());
			scopes.retainAll(clientDetails.getScope());
		}
		boolean upScopingHappened = false;
		if (!scopes.isEmpty()) {
			if (client.getScope() != null) {
				upScopingHappened = scopes.retainAll(client.getScope());
			}
			if (subjectToken.getScope() != null) {
				upScopingHappened = upScopingHappened || scopes.retainAll(subjectToken.getScope());
			}
		}
		if (upScopingHappened) {
			throw new InvalidScopeException("Up-scoping is not allowed");
		}
		token.setScope(scopes);

		// make it expire if necessary
		if (client.getAccessTokenValiditySeconds() != null && client.getAccessTokenValiditySeconds() > 0) {
			Date expiration = new Date(System.currentTimeMillis() + (client.getAccessTokenValiditySeconds() * 1000L));
			token.setExpiration(expiration);
		}

		AuthenticationHolderEntity authenticationHolder = subjectToken.getAuthenticationHolder();
		authenticationHolder.setId(null);
		authenticationHolder.setClientId(client.getClientId());
		authenticationHolder.setRedirectUri(null);
		authenticationHolder.setExtensions(new HashMap<>());
		authenticationHolder.setRequestParameters(tokenRequest.getRequestParameters());
		authenticationHolder.setResponseTypes(new HashSet<>());
		authenticationHolder.setScope(scopes);

		authenticationHolder = authenticationHolderRepository.save(authenticationHolder);
		token.setAuthenticationHolder(authenticationHolder);

		//Add approved site reference, if any
		OAuth2Request originalAuthRequest = subjectToken.getAuthenticationHolder().getAuthentication().getOAuth2Request();
		if (originalAuthRequest.getExtensions() != null && originalAuthRequest.getExtensions().containsKey("approved_site")) {
			Long apId = Long.parseLong((String) originalAuthRequest.getExtensions().get("approved_site"));
			ApprovedSite ap = approvedSiteService.getById(apId);
			token.setApprovedSite(ap);
		}

		//TODO: consider actor token?
		Set<String> audiences = new HashSet<>();
		if (StringUtils.hasText(tokenRequest.getRequestParameters().getOrDefault(ConnectRequestParameters.RESOURCE, null))) {
			audiences.addAll(
					Arrays.asList(
							tokenRequest.getRequestParameters().get(ConnectRequestParameters.RESOURCE).split(";")
					)
			);
		}
		if (StringUtils.hasText(tokenRequest.getRequestParameters().getOrDefault(ConnectRequestParameters.AUDIENCE, null))) {
			audiences.addAll(
					Arrays.asList(
							tokenRequest.getRequestParameters().get(ConnectRequestParameters.AUDIENCE).split(";")
					)
			);
		}
		if (!audiences.isEmpty()) {
			if (!clientDetailsEntityService.checkResourceIdsAreAllowedForClient(client.getClientId(), audiences)) {
				Set<String> mismatchedResIds = clientDetailsEntityService.getMismatchedResourceIds(
						client.getClientId(), originalAuthRequest.getResourceIds()
				);
				throw new InvalidTargetException("Identifier(s) '" + String.join(", ", mismatchedResIds) + "' not identifying any allowed resource for this client");
			}
		} else {
			audiences.add(client.getClientId());
		}

		// attach a refresh token, if this client is allowed to request them and the user gets the offline scope
		if (token.getScope().contains(OFFLINE_ACCESS)) {
			if (client.isAllowRefresh()) {
				OAuth2RefreshTokenEntity savedRefreshToken = createRefreshToken(client, token.getAuthenticationHolder(), audiences);
				token.setRefreshToken(savedRefreshToken);
			} else {
				throw new InvalidScopeException("Not authorized to request " + OFFLINE_ACCESS);
			}
		}

		JWTClaimsSet originalJwtClaims;
		try {
			originalJwtClaims = subjectToken.getJwtValue().getJWTClaimsSet();
		} catch (ParseException e) {
			throw new InvalidTokenException("Provided token could not be parsed as valid JWT");
		}

		audiences.add(client.getClientId());
		JWTClaimsSet.Builder builder = new JWTClaimsSet.Builder()
				.issuer(config.getConfigBean().getIssuer())
				.expirationTime(token.getExpiration())
				.audience(new ArrayList<>(audiences))
				.subject(originalJwtClaims.getSubject())
				.claim(CLIENT_ID, client.getClientId())
				.issueTime(new Date())
				.jwtID(UUID.randomUUID().toString())
				.claim(SCOPE, Joiner.on(SCOPE_SEPARATOR).join(token.getScope()));
		if (originalJwtClaims.getClaim(AUTH_TIME) != null) {
			builder = builder.claim(AUTH_TIME, originalJwtClaims.getClaim(AUTH_TIME));
		}
		if (originalJwtClaims.getClaim(ACR) != null) {
			builder = builder.claim(ACR, originalJwtClaims.getClaim(ACR));
		}
		if (originalJwtClaims.getClaim(AMR) != null) {
			builder = builder.claim(AMR, originalJwtClaims.getClaim(AMR));
		}

		JWTClaimsSet claims = builder.build();

		JWSAlgorithm signingAlg = jwtService.getDefaultSigningAlgorithm();
		JWSHeader header = new JWSHeader(signingAlg, JOSEObjectType.JWT, null, null, null, null, null, null, null, null,
				jwtService.getDefaultSignerKeyId(), true, null, null);
		SignedJWT signed = new SignedJWT(header, claims);

		jwtService.signJwt(signed);
		token.setJwtValue(signed);

		token.getAdditionalInformation().put(ISSUED_TOKEN_TYPE, requestedTokenType);
		OAuth2AccessTokenEntity savedToken = tokenServices.saveAccessToken(token);
		if (savedToken.getRefreshToken() != null) {
			tokenServices.saveRefreshToken(savedToken.getRefreshToken());
		}

		return savedToken;
	}

	@Override
	protected boolean isSupported() {
		return true;
	}

	private OAuth2RefreshTokenEntity createRefreshToken(
			ClientDetailsEntity client,
			AuthenticationHolderEntity authHolder,
			Set<String> resources
	) {
		OAuth2RefreshTokenEntity refreshToken = new OAuth2RefreshTokenEntity();
		JWTClaimsSet.Builder refreshClaims = new JWTClaimsSet.Builder();

		// make it expire if necessary
		if (client.getRefreshTokenValiditySeconds() != null) {
			Date expiration = new Date(System.currentTimeMillis() + (client.getRefreshTokenValiditySeconds() * 1000L));
			refreshToken.setExpiration(expiration);
			refreshClaims.expirationTime(expiration);
		}

		// set a random identifier
		refreshClaims.jwtID(UUID.randomUUID().toString());
		refreshClaims.issuer(config.getConfigBean().getIssuer());

		if (resources == null || resources.isEmpty()) {
			String audience = client.getClientId();
			if (!Strings.isNullOrEmpty(audience)) {
				refreshClaims.audience(Lists.newArrayList(audience));
			}
		} else {
			refreshClaims.audience(Lists.newArrayList(resources));
		}


		JWTClaimsSet claims = refreshClaims.build();
		JWSAlgorithm signingAlg = jwtService.getDefaultSigningAlgorithm();
		JWSHeader header = new JWSHeader(signingAlg, JOSEObjectType.JWT, null, null, null, null, null, null, null, null,
				jwtService.getDefaultSignerKeyId(), true, null, null);
		SignedJWT signed = new SignedJWT(header, claims);

		jwtService.signJwt(signed);
		refreshToken.setJwt(signed);

		//Add the authentication
		refreshToken.setAuthenticationHolder(authHolder);
		refreshToken.setClient(client);

		// save the token first so that we can set it to a member of the access token (NOTE: is this step necessary?)
		return getTokenServices().saveRefreshToken(refreshToken);
	}

	@Override
	public String toString() {
		return "OAuth2TokenExchangeGranter{" +
				"supportedSubjectTokenTypes=" + supportedSubjectTokenTypes +
				", supportedRequestedTokenTypes=" + supportedRequestedTokenTypes +
				", supportedActorTokenTypes=" + supportedActorTokenTypes +
				", supported=" + isSupported() +
				'}';
	}

}