package study.reactivesecurity;

import org.springframework.boot.SpringBootConfiguration;
import org.springframework.boot.autoconfigure.EnableAutoConfiguration;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Import;
import org.springframework.security.config.web.server.ServerHttpSecurity;
import org.springframework.security.core.userdetails.MapReactiveUserDetailsService;
import org.springframework.security.core.userdetails.User;
import org.springframework.security.crypto.bcrypt.BCryptPasswordEncoder;
import org.springframework.security.crypto.password.PasswordEncoder;
import org.springframework.security.authentication.UserDetailsRepositoryReactiveAuthenticationManager;
import org.springframework.security.config.Customizer;
import org.springframework.security.web.server.context.NoOpServerSecurityContextRepository;
import org.springframework.security.web.server.SecurityWebFilterChain;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.RestController;
import reactor.core.publisher.Mono;
import reactor.core.scheduler.Schedulers;
import org.springframework.security.core.context.ReactiveSecurityContextHolder;
import org.springframework.security.core.context.SecurityContextHolder;
import org.springframework.security.config.annotation.method.configuration.EnableReactiveMethodSecurity;
import org.springframework.security.access.prepost.PreAuthorize;
import org.springframework.web.bind.annotation.PathVariable;
import org.springframework.messaging.handler.annotation.MessageMapping;
import org.springframework.stereotype.Controller;
import org.springframework.security.config.annotation.rsocket.RSocketSecurity;
import org.springframework.security.rsocket.core.PayloadSocketAcceptorInterceptor;
import org.springframework.core.annotation.Order;
import org.springframework.security.web.server.util.matcher.ServerWebExchangeMatchers;
import org.springframework.security.oauth2.jwt.NimbusReactiveJwtDecoder;
import org.springframework.security.oauth2.jwt.JwtValidators;
import org.springframework.security.oauth2.jwt.JwtClaimValidator;
import org.springframework.security.oauth2.core.DelegatingOAuth2TokenValidator;
import org.springframework.security.oauth2.jose.jws.SignatureAlgorithm;
import org.springframework.security.web.server.context.WebSessionServerSecurityContextRepository;
import org.springframework.security.web.server.savedrequest.NoOpServerRequestCache;
import org.springframework.security.web.server.authentication.HttpStatusServerEntryPoint;
import org.springframework.security.web.server.authentication.logout.HttpStatusReturningServerLogoutSuccessHandler;
import org.springframework.security.web.server.csrf.CsrfToken;
import org.springframework.http.HttpStatus;
import org.springframework.http.HttpMethod;
import org.springframework.web.server.ServerWebExchange;
import org.springframework.web.bind.annotation.PostMapping;

@SpringBootConfiguration
@EnableAutoConfiguration
@Import({SecurityLab.Api.class, SecurityLab.Orders.class, SecurityLab.Messages.class})
@EnableReactiveMethodSecurity
public class SecurityLab {
    @Bean(destroyMethod="close") JwtKeys jwtKeys() throws Exception { return new JwtKeys(); }
    @Bean @Order(1) SecurityWebFilterChain jwtChain(ServerHttpSecurity http, JwtKeys keys) {
        var decoder=NimbusReactiveJwtDecoder.withJwkSetUri(keys.issuer()+"/jwks")
            .jwsAlgorithm(SignatureAlgorithm.RS256).build();
        decoder.setJwtValidator(new DelegatingOAuth2TokenValidator<>(JwtValidators.createDefaultWithIssuer(keys.issuer()),
            new JwtClaimValidator<java.util.List<String>>("aud", aud -> aud!=null&&aud.contains("u31-api"))));
        return http.securityMatcher(ServerWebExchangeMatchers.pathMatchers("/jwt/**"))
            .securityContextRepository(NoOpServerSecurityContextRepository.getInstance())
            .oauth2ResourceServer(oauth -> oauth.jwt(jwt -> jwt.jwtDecoder(decoder)))
            .authorizeExchange(a -> a.anyExchange().hasAuthority("SCOPE_read")).build();
    }
    @Bean PayloadSocketAcceptorInterceptor payloadSecurity(RSocketSecurity socket) {
        return socket.simpleAuthentication(Customizer.withDefaults())
            .authorizePayload(a -> a.route("admin").hasRole("ADMIN")
                .anyRequest().authenticated().anyExchange().permitAll()).build();
    }
    @Bean PasswordEncoder passwords() { return new BCryptPasswordEncoder(6); }
    @Bean MapReactiveUserDetailsService users(PasswordEncoder passwords) {
        return new MapReactiveUserDetailsService(
            User.withUsername("alice").password(passwords.encode("lab-password")).roles("READER").build(),
            User.withUsername("bob").password(passwords.encode("lab-password")).roles("READER").build(),
            User.withUsername("admin").password(passwords.encode("lab-password")).roles("READER","ADMIN").build());
    }
    @Bean UserDetailsRepositoryReactiveAuthenticationManager authenticationManager(
            MapReactiveUserDetailsService users, PasswordEncoder passwords) {
        var manager=new UserDetailsRepositoryReactiveAuthenticationManager(users);
        manager.setPasswordEncoder(passwords);
        return manager;
    }
    @Bean @Order(3) SecurityWebFilterChain apiChain(ServerHttpSecurity http) {
        return http.securityContextRepository(NoOpServerSecurityContextRepository.getInstance())
            .httpBasic(Customizer.withDefaults())
            .authorizeExchange(a -> a.pathMatchers("/api/admin").hasRole("ADMIN")
                .anyExchange().authenticated()).build();
    }
    @Bean @Order(2) SecurityWebFilterChain browserChain(ServerHttpSecurity http) {
        var repository=new WebSessionServerSecurityContextRepository();
        return http.securityMatcher(ServerWebExchangeMatchers.pathMatchers("/browser/**"))
            .securityContextRepository(repository)
            .requestCache(cache -> cache.requestCache(NoOpServerRequestCache.getInstance()))
            .httpBasic(basic -> basic.securityContextRepository(repository))
            .exceptionHandling(e -> e.authenticationEntryPoint(new HttpStatusServerEntryPoint(HttpStatus.UNAUTHORIZED)))
            .logout(logout -> logout.requiresLogout(ServerWebExchangeMatchers.pathMatchers(HttpMethod.POST,"/browser/logout"))
                .logoutSuccessHandler(new HttpStatusReturningServerLogoutSuccessHandler(HttpStatus.OK)))
            .authorizeExchange(a -> a.pathMatchers("/browser/csrf").permitAll().anyExchange().authenticated()).build();
    }
    @RestController
    static class Api {
        private final Orders orders;
        Api(Orders orders) { this.orders=orders; }
        @GetMapping("/api/admin") Mono<String> admin() { return Mono.just("admin-ok"); }
        @GetMapping("/api/method") Mono<String> method() { return orders.approve(); }
        @GetMapping("/jwt/who") Mono<String> jwt() {
            return ReactiveSecurityContextHolder.getContext().map(c -> "jwt:"+c.getAuthentication().getName());
        }
        @GetMapping("/browser/csrf") Mono<String> csrf(ServerWebExchange exchange) {
            Mono<CsrfToken> token=exchange.getAttribute(CsrfToken.class.getName());
            return token.map(CsrfToken::getToken);
        }
        @GetMapping("/browser/who") Mono<String> browserWho() {
            return ReactiveSecurityContextHolder.getContext().map(c -> "session:"+c.getAuthentication().getName());
        }
        @PostMapping("/browser/write") Mono<String> write() { return Mono.just("write-ok"); }
        @GetMapping("/api/orders/{owner}") Mono<String> order(@PathVariable String owner) {
            return orders.read(owner);
        }
        @GetMapping("/api/who") Mono<String> who() {
            return Mono.just("switch").publishOn(Schedulers.boundedElastic()).flatMap(ignored ->
                ReactiveSecurityContextHolder.getContext().map(context -> {
                    var local=SecurityContextHolder.getContext().getAuthentication();
                    return "ctx="+context.getAuthentication().getName()+"|threadLocal="+
                        (local==null?"none":local.getName())+"@"+Thread.currentThread().getName();
                }));
        }
    }
    static class Orders {
        @PreAuthorize("#owner == authentication.name")
        public Mono<String> read(String owner) { return Mono.just("order-of:"+owner); }
        @PreAuthorize("hasRole('ADMIN')")
        public Mono<String> approve() { return Mono.just("approved"); }
    }
    @Controller
    static class Messages {
        @MessageMapping("read") Mono<String> read(String input) {
            return ReactiveSecurityContextHolder.getContext().map(c -> "rsocket:"+c.getAuthentication().getName());
        }
        @MessageMapping("admin") Mono<String> admin(String input) { return Mono.just("rsocket-admin"); }
    }
}