package mx.gob.slp.baseadminweb.configuration.security;

import org.springframework.boot.autoconfigure.condition.ConditionalOnProperty;
import org.springframework.boot.context.properties.EnableConfigurationProperties;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.security.core.GrantedAuthority;
import org.springframework.security.core.authority.SimpleGrantedAuthority;
import org.springframework.security.oauth2.jose.jws.MacAlgorithm;
import org.springframework.security.oauth2.jwt.Jwt;
import org.springframework.security.oauth2.jwt.JwtDecoder;
import org.springframework.security.oauth2.jwt.JwtEncoder;
import org.springframework.security.oauth2.jwt.NimbusJwtDecoder;
import org.springframework.security.oauth2.jwt.NimbusJwtEncoder;
import org.springframework.security.oauth2.server.resource.authentication.JwtAuthenticationConverter;
import org.springframework.security.oauth2.server.resource.authentication.JwtGrantedAuthoritiesConverter;

import com.nimbusds.jose.JWSAlgorithm;
import com.nimbusds.jose.jwk.JWKSet;
import com.nimbusds.jose.jwk.OctetSequenceKey;
import com.nimbusds.jose.jwk.source.ImmutableJWKSet;

import javax.crypto.SecretKey;
import javax.crypto.spec.SecretKeySpec;
import java.nio.charset.StandardCharsets;
import java.util.ArrayList;
import java.util.Collection;
import java.util.List;

@Configuration
@EnableConfigurationProperties({JwtProperties.class, CognitoProperties.class})
public class JwtConfig {

    @Bean
    @ConditionalOnProperty(prefix = "app.cognito", name = "enabled", havingValue = "false")
    JwtEncoder jwtEncoder(JwtProperties jwtProperties) {
        OctetSequenceKey jwk = new OctetSequenceKey.Builder(jwtProperties.getSecret().getBytes(StandardCharsets.UTF_8))
                .algorithm(JWSAlgorithm.HS256)
                .build();
        return new NimbusJwtEncoder(new ImmutableJWKSet<>(new JWKSet(jwk)));
    }

    @Bean
    @ConditionalOnProperty(prefix = "app.cognito", name = "enabled", havingValue = "false")
    JwtDecoder localJwtDecoder(JwtProperties jwtProperties) {
        return NimbusJwtDecoder.withSecretKey(secretKey(jwtProperties))
                .macAlgorithm(MacAlgorithm.HS256)
                .build();
    }

    @Bean
    @ConditionalOnProperty(prefix = "app.cognito", name = "enabled", havingValue = "true")
    JwtDecoder cognitoJwtDecoder(CognitoProperties cognitoProperties) {
        return NimbusJwtDecoder.withJwkSetUri(cognitoProperties.getJwkSetUri()).build();
    }

    @Bean
    JwtAuthenticationConverter jwtAuthenticationConverter(CognitoProperties cognitoProperties) {
        JwtGrantedAuthoritiesConverter localRolesConverter = new JwtGrantedAuthoritiesConverter();
        localRolesConverter.setAuthoritiesClaimName("roles");
        localRolesConverter.setAuthorityPrefix("");

        JwtAuthenticationConverter authenticationConverter = new JwtAuthenticationConverter();
        authenticationConverter.setJwtGrantedAuthoritiesConverter(jwt -> {
            Collection<GrantedAuthority> authorities = new ArrayList<>();
            List<String> groups = jwt.getClaimAsStringList("cognito:groups");
            CognitoRoleMapper.toRoles(groups, cognitoProperties).stream()
                    .map(SimpleGrantedAuthority::new)
                    .forEach(authorities::add);
            authorities.addAll(localRolesConverter.convert(jwt));
            return authorities;
        });
        return authenticationConverter;
    }

    private SecretKey secretKey(JwtProperties jwtProperties) {
        return new SecretKeySpec(jwtProperties.getSecret().getBytes(), "HmacSHA256");
    }
}