package com.rabbitmq.authorization_server;

import java.security.KeyPair;
import java.security.KeyPairGenerator;
import java.security.interfaces.RSAPrivateKey;
import java.security.interfaces.RSAPublicKey;
import java.util.UUID;
import java.util.Collection;
import java.util.List;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.core.annotation.Order;
import org.springframework.http.MediaType;
import org.springframework.security.authentication.AbstractAuthenticationToken;
import org.springframework.security.config.Customizer;
import org.springframework.security.config.annotation.web.builders.HttpSecurity;
import org.springframework.security.config.annotation.web.configuration.EnableWebSecurity;
import org.springframework.security.core.userdetails.UserDetailsService;
import org.springframework.security.oauth2.jwt.JwtDecoder;
import org.springframework.security.oauth2.server.authorization.OAuth2TokenType;
import org.springframework.security.oauth2.server.authorization.config.annotation.web.configuration.OAuth2AuthorizationServerConfiguration;
import org.springframework.security.oauth2.server.authorization.config.annotation.web.configurers.OAuth2AuthorizationServerConfigurer;
import org.springframework.security.oauth2.server.authorization.settings.AuthorizationServerSettings;
import org.springframework.security.oauth2.server.authorization.token.JwtEncodingContext;
import org.springframework.security.oauth2.server.authorization.token.OAuth2TokenCustomizer;
import org.springframework.security.provisioning.InMemoryUserDetailsManager;
import org.springframework.security.web.SecurityFilterChain;
import org.springframework.security.web.authentication.LoginUrlAuthenticationEntryPoint;
import org.springframework.security.web.util.matcher.MediaTypeRequestMatcher;
import org.springframework.security.oauth2.server.authorization.token.OAuth2TokenClaimsContext;

import org.springframework.security.oauth2.core.AuthorizationGrantType;

import com.nimbusds.jose.jwk.JWKSet;
import com.nimbusds.jose.jwk.RSAKey;
import com.nimbusds.jose.jwk.source.ImmutableJWKSet;
import com.nimbusds.jose.jwk.source.JWKSource;
import com.nimbusds.jose.proc.SecurityContext;

@Configuration
@EnableWebSecurity
public class SecurityConfig {
	
	@Bean 
	@Order(1)
	public SecurityFilterChain authorizationServerSecurityFilterChain(HttpSecurity http)
			throws Exception {
		OAuth2AuthorizationServerConfigurer authorizationServerConfigurer =
				OAuth2AuthorizationServerConfigurer.authorizationServer();

		http
			.securityMatcher(authorizationServerConfigurer.getEndpointsMatcher())
			.with(authorizationServerConfigurer, (authorizationServer) ->
				authorizationServer
					.oidc(Customizer.withDefaults())	// Enable OpenID Connect 1.0
			)
			.authorizeHttpRequests((authorize) ->
				authorize
					.anyRequest().authenticated()
			)
			// Redirect to the login page when not authenticated from the
			// authorization endpoint
			.exceptionHandling((exceptions) -> exceptions
				.defaultAuthenticationEntryPointFor(
					new LoginUrlAuthenticationEntryPoint("/login"),
					new MediaTypeRequestMatcher(MediaType.TEXT_HTML)
				)
			);

		return http.build();
	}

	@Bean 
	@Order(2)
	public SecurityFilterChain defaultSecurityFilterChain(HttpSecurity http)
			throws Exception {
		http
			.authorizeHttpRequests((authorize) -> authorize
				.anyRequest().authenticated()
			)
			// Form login handles the redirect to the login page from the
			// authorization server filter chain
			.formLogin(Customizer.withDefaults());

		return http.build();
	}

	@Bean 
	public UserDetailsService userDetailsService(UsersConfiguration users) {
		return new InMemoryUserDetailsManager(users.getUserDetails());
	}

	@Bean 
	public JWKSource<SecurityContext> jwkSource() {
		KeyPair keyPair = generateRsaKey();
		RSAPublicKey publicKey = (RSAPublicKey) keyPair.getPublic();
		RSAPrivateKey privateKey = (RSAPrivateKey) keyPair.getPrivate();
		RSAKey rsaKey = new RSAKey.Builder(publicKey)
				.privateKey(privateKey)
				.keyID(UUID.randomUUID().toString())
				.build();
		JWKSet jwkSet = new JWKSet(rsaKey);
		return new ImmutableJWKSet<>(jwkSet);
	}

	private static KeyPair generateRsaKey() { 
		KeyPair keyPair;
		try {
			KeyPairGenerator keyPairGenerator = KeyPairGenerator.getInstance("RSA");
			keyPairGenerator.initialize(2048);
			keyPair = keyPairGenerator.generateKeyPair();
		}
		catch (Exception ex) {
			throw new IllegalStateException(ex);
		}
		return keyPair;
	}

	Logger logger = LoggerFactory.getLogger(SecurityConfig.class);

	@Bean 
	public OAuth2TokenCustomizer<OAuth2TokenClaimsContext>  accessTokenCustomizer() {
		logger.info("Creating accessTokenCustomizer ...");
		return (context) -> {
			logger.info("Calling accessTokenCustomizer with tokenType: {}", context.getTokenType().getValue());
			AbstractAuthenticationToken principal = context.getPrincipal();
			logger.info("registered client: {}", context.getRegisteredClient());
			logger.info("principal : {}", principal);
			logger.info("token format : {} ",
					context.getRegisteredClient().getTokenSettings().getAccessTokenFormat().getValue());
			logger.info("authorities : {}", principal.getAuthorities());
			logger.info("authorized scopes : {}", context.getAuthorizedScopes());
				
			if (AuthorizationGrantType.CLIENT_CREDENTIALS.equals(context.getAuthorizationGrantType())) {
				Collection<String> extra_scope = context.getRegisteredClient().getScopes();
				logger.info("granting extra_scope: {}", extra_scope);
				context.getClaims()
					.claim("extra_scope", extra_scope);
			} else {
				Collection<String> extra_scope = ScopeAuthority.getAuthorites(principal);
				List<String> audience = AudienceAuthority.getAll(principal);
				logger.info("granting extra_scope: {}", extra_scope);
				logger.info("granting audience: {}", audience);
				context.getClaims()
					.audience(audience)
					.claim("extra_scope", extra_scope);
			}
		};
	}
	@Bean
	public OAuth2TokenCustomizer<JwtEncodingContext> jwtTokenCustomizer() {
		logger.info("Creating jwtTokenCustomizer ...");
		return (context) -> {
			logger.info("Calling jwtTokenCustomizer with tokenType: {}", context.getTokenType().getValue());
			if (OAuth2TokenType.ACCESS_TOKEN.equals(context.getTokenType())) {
				AbstractAuthenticationToken principal = context.getPrincipal();
				logger.info("registered client: {}", context.getRegisteredClient());
				logger.info("principal : {}", principal);
				logger.info("token format : {} ",
					context.getRegisteredClient().getTokenSettings().getAccessTokenFormat().getValue());
				logger.info("authorities : {}", principal.getAuthorities());
				logger.info("authorized scopes : {}", context.getAuthorizedScopes());
				
				context.getClaims()
					.audience(AudienceAuthority.getAll(principal))
					.claim("extra_scope", ScopeAuthority.getAuthorites(principal));				
			}
		};
	}


	@Bean 
	public JwtDecoder jwtDecoder(JWKSource<SecurityContext> jwkSource) {
		return OAuth2AuthorizationServerConfiguration.jwtDecoder(jwkSource);
	}

	@Bean 
	public AuthorizationServerSettings authorizationServerSettings() {
		return AuthorizationServerSettings.builder().build();
	}

}
