package study.oauthlab;

import java.time.Instant;
import java.time.Duration;
import java.util.*;
import java.util.concurrent.atomic.AtomicReference;
import com.nimbusds.jose.*;
import com.nimbusds.jose.crypto.*;
import com.nimbusds.jose.jwk.*;
import com.nimbusds.jose.jwk.gen.*;
import com.nimbusds.jose.jwk.source.JWKSource;
import com.nimbusds.jose.proc.SecurityContext;
import com.nimbusds.jwt.*;
import jakarta.servlet.FilterChain;
import jakarta.servlet.http.*;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.context.annotation.*;
import org.springframework.core.annotation.Order;
import org.springframework.http.MediaType;
import org.springframework.security.config.Customizer;
import org.springframework.security.config.annotation.web.builders.HttpSecurity;
import org.springframework.security.core.userdetails.*;
import org.springframework.security.provisioning.InMemoryUserDetailsManager;
import org.springframework.security.oauth2.core.*;
import org.springframework.security.oauth2.core.oidc.OidcScopes;
import org.springframework.security.oauth2.server.authorization.client.*;
import org.springframework.security.oauth2.server.authorization.settings.*;
import org.springframework.security.oauth2.server.authorization.token.*;
import org.springframework.security.oauth2.server.authorization.OAuth2TokenType;
import org.springframework.security.web.*;
import org.springframework.security.web.authentication.LoginUrlAuthenticationEntryPoint;
import org.springframework.security.web.context.SecurityContextHolderFilter;
import org.springframework.security.web.csrf.CsrfToken;
import org.springframework.security.web.util.matcher.MediaTypeRequestMatcher;
import org.springframework.web.bind.annotation.*;
import org.springframework.web.filter.OncePerRequestFilter;

@Configuration
@Profile("as")
class AuthorizationServer {
    @Bean Keys keys() throws Exception { return new Keys(); }
    @Bean JWKSource<SecurityContext> jwkSource(Keys keys) {
        return (selector, context) -> selector.select(new JWKSet(keys.published()));
    }
    @Bean RegisteredClientRepository clients(@Value("${lab.client}") String client) {
        var interactive = RegisteredClient.withId("synthetic-client")
            .clientId("lab-client").clientSecret("{noop}lab-secret")
            .clientAuthenticationMethod(ClientAuthenticationMethod.CLIENT_SECRET_BASIC)
            .authorizationGrantType(AuthorizationGrantType.AUTHORIZATION_CODE)
            .authorizationGrantType(AuthorizationGrantType.REFRESH_TOKEN)
            .redirectUri(client + "/login/oauth2/code/lab")
            .redirectUri("http://127.0.0.1/callback")
            .scope(OidcScopes.OPENID).scope(OidcScopes.PROFILE).scope("orders.read").scope("orders.write")
            .clientSettings(ClientSettings.builder().requireProofKey(true).requireAuthorizationConsent(false).build())
            .tokenSettings(TokenSettings.builder().accessTokenTimeToLive(Duration.ofMinutes(5))
                .reuseRefreshTokens(false).build()).build();
        var reference = RegisteredClient.withId("synthetic-reference").clientId("reference-client").clientSecret("{noop}reference-secret")
            .clientAuthenticationMethod(ClientAuthenticationMethod.CLIENT_SECRET_BASIC)
            .authorizationGrantType(AuthorizationGrantType.CLIENT_CREDENTIALS).scope("orders.read")
            .tokenSettings(TokenSettings.builder().accessTokenFormat(OAuth2TokenFormat.REFERENCE).build()).build();
        return new InMemoryRegisteredClientRepository(interactive, reference);
    }
    @Bean AuthorizationServerSettings settings(@Value("${lab.issuer}") String issuer) {
        return AuthorizationServerSettings.builder().issuer(issuer).build();
    }
    @Bean UserDetailsService users() {
        return new InMemoryUserDetailsManager(User.withUsername("alice").password("{noop}lab-password").roles("USER").build());
    }
    @Bean OAuth2TokenCustomizer<JwtEncodingContext> claims(Keys keys) {
        return context -> {
            context.getJwsHeader().keyId(keys.active.get());
            if (OAuth2TokenType.ACCESS_TOKEN.equals(context.getTokenType()))
                context.getClaims().audience(List.of("orders-api")).claim("tenant", "t1");
            if ("id_token".equals(context.getTokenType().getValue()) && keys.badNonce)
                context.getClaims().claim("nonce", "deliberately-wrong-synthetic-nonce");
        };
    }
    @Bean OAuth2TokenCustomizer<OAuth2TokenClaimsContext> referenceClaims() {
        return context -> context.getClaims().audience(List.of("orders-api")).claim("tenant", "t1");
    }
    @Bean @Order(1) SecurityFilterChain protocol(HttpSecurity http, Keys keys) throws Exception {
        http.oauth2AuthorizationServer(as -> { http.securityMatcher(as.getEndpointsMatcher()); as.oidc(Customizer.withDefaults()); })
            .authorizeHttpRequests(a -> a.anyRequest().authenticated())
            .exceptionHandling(e -> e.defaultAuthenticationEntryPointFor(new LoginUrlAuthenticationEntryPoint("/login"), new MediaTypeRequestMatcher(MediaType.TEXT_HTML)));
        http.addFilterBefore(new OncePerRequestFilter() {
            protected void doFilterInternal(HttpServletRequest r, HttpServletResponse s, FilterChain c) throws java.io.IOException, jakarta.servlet.ServletException {
                if ((r.getRequestURI().equals("/oauth2/jwks") && keys.jwksDown) ||
                    (r.getRequestURI().equals("/oauth2/introspect") && keys.introspectionDown)) { s.setStatus(503); return; }
                c.doFilter(r, s);
            }
        }, SecurityContextHolderFilter.class);
        return http.build();
    }
    @Bean @Order(2) SecurityFilterChain browser(HttpSecurity http) throws Exception {
        http.authorizeHttpRequests(a -> a.requestMatchers("/lab/**").permitAll().anyRequest().authenticated())
            .csrf(c -> c.ignoringRequestMatchers("/lab/control", "/lab/mint"))
            .formLogin(Customizer.withDefaults());
        return http.build();
    }
    static class Keys {
        final RSAKey a; final RSAKey b; final RSAKey foreign;
        final AtomicReference<String> active = new AtomicReference<>("A");
        volatile String publication = "A";
        volatile boolean jwksDown, introspectionDown, badNonce;
        Keys() throws Exception {
            a = new RSAKeyGenerator(2048).keyID("A").generate();
            b = new RSAKeyGenerator(2048).keyID("B").generate();
            foreign = new RSAKeyGenerator(2048).keyID("foreign").generate();
        }
        List<JWK> published() { return switch (publication) { case "AB" -> List.of(a, b); case "B" -> List.of(b); default -> List.of(a); }; }
    }
    @RestController @Profile("as") static class Controls {
        final Keys keys; final String control; final String issuer;
        Controls(Keys k, @Value("${LAB_CONTROL}") String c, @Value("${lab.issuer}") String i) {keys=k; control=c; issuer=i;}
        void check(String value) { if (!java.security.MessageDigest.isEqual(control.getBytes(java.nio.charset.StandardCharsets.UTF_8), value.getBytes(java.nio.charset.StandardCharsets.UTF_8))) throw new org.springframework.web.server.ResponseStatusException(org.springframework.http.HttpStatus.FORBIDDEN); }
        @GetMapping("/lab/csrf") Map<String,String> csrf(CsrfToken t) { return Map.of("token",t.getToken(),"parameterName",t.getParameterName()); }
        @PostMapping("/lab/control") Map<String,Object> control(@RequestHeader("X-Lab-Control") String secret, @RequestParam Map<String,String> p) {
            check(secret);
            if (p.containsKey("publish")) keys.publication=p.get("publish");
            if (p.containsKey("active")) keys.active.set(p.get("active"));
            if (p.containsKey("jwksDown")) keys.jwksDown=Boolean.parseBoolean(p.get("jwksDown"));
            if (p.containsKey("introspectionDown")) keys.introspectionDown=Boolean.parseBoolean(p.get("introspectionDown"));
            if (p.containsKey("badNonce")) keys.badNonce=Boolean.parseBoolean(p.get("badNonce"));
            return Map.of("publication", keys.publication, "active", keys.active.get());
        }
        @PostMapping("/lab/mint") Map<String,String> mint(@RequestHeader("X-Lab-Control") String secret, @RequestParam(defaultValue="valid") String kind) throws Exception {
            check(secret);
            Instant now=Instant.now();
            var c=new JWTClaimsSet.Builder().issuer(kind.equals("issuer")?"https://wrong.invalid":issuer)
                .subject("alice").audience(kind.equals("audience")?"other-api":"orders-api")
                .issueTime(Date.from(now)).claim("tenant", "t1").claim("scope",kind.equals("scope")?"profile":"orders.read orders.write")
                .expirationTime(Date.from(now.plusSeconds(kind.equals("expired")?-120:300)));
            if(kind.equals("no-exp")) c.expirationTime(null);
            if(kind.equals("future")) c.notBeforeTime(Date.from(now.plusSeconds(300)));
            if(kind.equals("no-tenant")) c.claim("tenant",null);
            if(kind.equals("tenant")) c.claim("tenant","t2");
            String kid = kind.equals("unknown-kid")?"unknown":keys.active.get();
            JWSAlgorithm alg=kind.equals("algorithm")?JWSAlgorithm.RS512:JWSAlgorithm.RS256;
            SignedJWT token=new SignedJWT(new JWSHeader.Builder(alg).type(JOSEObjectType.JWT).keyID(kid).build(),c.build());
            RSAKey signer=kind.equals("signature")?keys.foreign:(keys.active.get().equals("A")?keys.a:keys.b);
            token.sign(new RSASSASigner(signer));
            return Map.of("token",token.serialize());
        }
    }
}