package com.example.apigateway.config;

import java.util.HashSet;
import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.stream.Collectors;

import org.springframework.beans.factory.annotation.Value;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.security.config.annotation.web.reactive.EnableWebFluxSecurity;
import org.springframework.security.core.GrantedAuthority;
import org.springframework.security.core.authority.SimpleGrantedAuthority;
import org.springframework.security.core.authority.mapping.GrantedAuthoritiesMapper;
import org.springframework.security.oauth2.client.oidc.userinfo.OidcReactiveOAuth2UserService;
import org.springframework.security.oauth2.client.oidc.userinfo.OidcUserRequest;
import org.springframework.security.oauth2.client.registration.ClientRegistration.ProviderDetails;
import org.springframework.security.oauth2.client.userinfo.ReactiveOAuth2UserService;
import org.springframework.security.oauth2.core.OAuth2AccessToken;
import org.springframework.security.oauth2.core.oidc.OidcIdToken;
import org.springframework.security.oauth2.core.oidc.OidcUserInfo;
import org.springframework.security.oauth2.core.oidc.user.DefaultOidcUser;
import org.springframework.security.oauth2.core.oidc.user.OidcUser;
import org.springframework.security.oauth2.core.oidc.user.OidcUserAuthority;
import org.springframework.security.oauth2.core.user.DefaultOAuth2User;
import org.springframework.security.oauth2.core.user.OAuth2UserAuthority;
import org.springframework.util.StringUtils;
import org.springframework.web.reactive.function.client.WebClient;

import com.example.apigateway.security.dto.UserResonseDto;

import reactor.core.publisher.Mono;

/**
 * Ref:
 * https://docs.spring.io/spring-security/reference/reactive/oauth2/login/advanced.html
 */
@Configuration
@EnableWebFluxSecurity
public class SecurityConfig {

    private final WebClient webClient = WebClient.create(); // WebClient configured to point to your backend service

    @Value("${app.userRolesUrl}")
    private String userRolesUrl;

    @Bean
    public ReactiveOAuth2UserService<OidcUserRequest, OidcUser> oidcUserService() {
        final OidcReactiveOAuth2UserService delegate = new OidcReactiveOAuth2UserService();

        return (userRequest) -> {
            // Delegate to the default implementation for loading a user
            return delegate.loadUser(userRequest)
                    .flatMap((oidcUser) -> {
                        String username = oidcUser.getAttribute("sub");

                        return fetchRolesFromBackend(username)
                                .map(roles -> {
                                    Set<GrantedAuthority> authorities = roles.stream()
                                            .map(role -> new SimpleGrantedAuthority("ROLE_" + role.toUpperCase()))
                                            .collect(Collectors.toSet());

                                    // Keep original identity scopes if needed
                                    authorities.addAll(oidcUser.getAuthorities());

                                    ProviderDetails providerDetails = userRequest.getClientRegistration()
                                            .getProviderDetails();
                                    String userNameAttributeName = providerDetails.getUserInfoEndpoint()
                                            .getUserNameAttributeName();
                                    if (StringUtils.hasText(userNameAttributeName)) {
                                        return new DefaultOidcUser(authorities, oidcUser.getIdToken(),
                                                oidcUser.getUserInfo(), userNameAttributeName);
                                    } else {
                                        return new DefaultOidcUser(authorities, oidcUser.getIdToken(),
                                                oidcUser.getUserInfo());
                                    }
                                });
                    });
        };
    }

    private Mono<List<String>> fetchRolesFromBackend(String username) {
        Mono<List<String>> response = this.webClient.get()
                .uri(userRolesUrl, username)
                .retrieve()
                .bodyToMono(UserResonseDto.class)
                .map(UserResonseDto::getRoles)
                .onErrorReturn(List.of("")); // Can return a String indicating a fallback role on backend failure

        return response;
    }
}