package com.baiyue.common.security;

import com.nimbusds.jose.JWSAlgorithm;
import com.nimbusds.jose.JWSHeader;
import com.nimbusds.jose.crypto.RSASSASigner;
import com.nimbusds.jose.jwk.JWKSet;
import com.nimbusds.jose.jwk.RSAKey;
import com.nimbusds.jose.jwk.gen.RSAKeyGenerator;
import com.nimbusds.jwt.JWTClaimsSet;
import com.nimbusds.jwt.SignedJWT;
import com.sun.net.httpserver.HttpServer;
import org.junit.jupiter.api.Test;
import org.springframework.security.oauth2.jwt.JwtException;

import java.net.InetSocketAddress;
import java.net.ServerSocket;
import java.nio.charset.StandardCharsets;
import java.time.Instant;
import java.util.Date;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.concurrent.atomic.AtomicReference;

import static org.assertj.core.api.Assertions.assertThat;
import static org.assertj.core.api.Assertions.assertThatCode;
import static org.assertj.core.api.Assertions.assertThatThrownBy;

class JwksJwtDecoderAutoConfigurationTest {
    @Test
    void createsDecoderWhenAuthIsUnavailableButRejectsTokenVerification() throws Exception {
        int port;
        try (var socket = new ServerSocket(0)) {
            port = socket.getLocalPort();
        }
        var decoder = new JwksJwtDecoderAutoConfiguration().jwksJwtDecoder(
                "http://127.0.0.1:" + port + "/jwks", "auth-service", "lastleg-api");
        var key = new RSAKeyGenerator(2048).keyID("current").generate();
        String token = token(key, "auth-service", "lastleg-api", 120, JWSAlgorithm.RS256);
        assertThatThrownBy(() -> decoder.decode(token)).isInstanceOf(JwtException.class);
    }

    @Test
    void fetchesOnFirstVerificationCachesKeysAndRefreshesForRotation() throws Exception {
        var current = new RSAKeyGenerator(2048).keyID("current").generate();
        var next = new RSAKeyGenerator(2048).keyID("next").generate();
        var published = new AtomicReference<>(new JWKSet(current.toPublicJWK()));
        var requests = new AtomicInteger();
        var server = HttpServer.create(new InetSocketAddress("127.0.0.1", 0), 0);
        server.createContext("/jwks", exchange -> {
            requests.incrementAndGet();
            byte[] body = published.get().toString().getBytes(StandardCharsets.UTF_8);
            exchange.getResponseHeaders().add("Content-Type", "application/json");
            exchange.sendResponseHeaders(200, body.length);
            try (var output = exchange.getResponseBody()) {
                output.write(body);
            }
        });
        server.start();
        try {
            var decoder = new JwksJwtDecoderAutoConfiguration().jwksJwtDecoder(
                    "http://127.0.0.1:" + server.getAddress().getPort() + "/jwks",
                    "auth-service", "lastleg-api");
            assertThat(requests.get()).isZero();
            String valid = token(current, "auth-service", "lastleg-api", 120, JWSAlgorithm.RS256);
            assertThat(decoder.decode(valid).getSubject()).isEqualTo("42");
            assertThatCode(() -> decoder.decode(valid)).doesNotThrowAnyException();
            assertThat(requests.get()).isEqualTo(1);

            for (String invalid : new String[]{
                    token(current, "other", "lastleg-api", 120, JWSAlgorithm.RS256),
                    token(current, "auth-service", "other", 120, JWSAlgorithm.RS256),
                    token(current, "auth-service", "lastleg-api", -120, JWSAlgorithm.RS256),
                    token(current, "auth-service", "lastleg-api", 120, JWSAlgorithm.RS512)}) {
                assertThatThrownBy(() -> decoder.decode(invalid)).isInstanceOf(JwtException.class);
            }

            published.set(new JWKSet(next.toPublicJWK()));
            assertThat(decoder.decode(token(next, "auth-service", "lastleg-api", 120,
                    JWSAlgorithm.RS256)).getSubject()).isEqualTo("42");
            assertThat(requests.get()).isEqualTo(2);
        } finally {
            server.stop(0);
        }
    }

    private String token(RSAKey key, String issuer, String audience, int ttl, JWSAlgorithm algorithm)
            throws Exception {
        var jwt = new SignedJWT(new JWSHeader.Builder(algorithm).keyID(key.getKeyID()).build(),
                new JWTClaimsSet.Builder().subject("42").issuer(issuer).audience(audience)
                        .issueTime(Date.from(Instant.now().minusSeconds(300)))
                        .expirationTime(Date.from(Instant.now().plusSeconds(ttl))).build());
        jwt.sign(new RSASSASigner(key));
        return jwt.serialize();
    }
}