package com.baiyue.auth.config;

import com.nimbusds.jose.JWSAlgorithm;
import com.nimbusds.jose.jwk.JWKSet;
import com.nimbusds.jose.jwk.gen.RSAKeyGenerator;
import org.junit.jupiter.api.Test;
import org.springframework.security.oauth2.jose.jws.SignatureAlgorithm;
import org.springframework.security.oauth2.jwt.*;
import java.time.Instant;
import java.util.List;
import static org.assertj.core.api.Assertions.*;

class UnifiedJwtConfigurationTest {
    @Test
    void signsWithCurrentKeyAndRetainsPreviousPublicKeyForRotation() throws Exception {
        var current = new RSAKeyGenerator(2048).keyID("current").algorithm(JWSAlgorithm.RS256).generate();
        var previous = new RSAKeyGenerator(2048).keyID("previous").algorithm(JWSAlgorithm.RS256).generate();
        var properties = new UnifiedJwtProperties(new JWKSet(List.of(current, previous.toPublicJWK())).toString(false), "current");
        var encoder = new UnifiedJwtConfiguration().jwtEncoder(properties);
        var now = Instant.now();
        var token = encoder.encode(JwtEncoderParameters.from(
                JwsHeader.with(SignatureAlgorithm.RS256).keyId("current").build(),
                JwtClaimsSet.builder().subject("42").issuedAt(now).expiresAt(now.plusSeconds(60)).build()));
        var decoded = NimbusJwtDecoder.withPublicKey(current.toRSAPublicKey()).build().decode(token.getTokenValue());
        assertThat(decoded.getSubject()).isEqualTo("42");
        assertThat(decoded.getHeaders()).containsEntry("kid", "current").containsEntry("alg", "RS256");
        assertThatThrownBy(() -> NimbusJwtDecoder.withPublicKey(previous.toRSAPublicKey()).build().decode(token.getTokenValue()))
                .isInstanceOf(JwtException.class);
    }

    @Test
    void rejectsMissingPrivateKeyAndPreviousPrivateKey() throws Exception {
        var current = new RSAKeyGenerator(2048).keyID("current").generate();
        var previous = new RSAKeyGenerator(2048).keyID("previous").generate();
        var config = new UnifiedJwtConfiguration();
        assertThatThrownBy(() -> config.jwtEncoder(new UnifiedJwtProperties(new JWKSet(current.toPublicJWK()).toString(), "current")))
                .isInstanceOf(IllegalStateException.class);
        assertThatThrownBy(() -> config.jwtEncoder(new UnifiedJwtProperties(new JWKSet(List.of(current, previous)).toString(false), "current")))
                .isInstanceOf(IllegalStateException.class);
        assertThatThrownBy(() -> config.jwtEncoder(new UnifiedJwtProperties("", "current")))
                .isInstanceOf(IllegalStateException.class);
    }
}