package mx.gob.slp.baseadminweb.configuration.security;

import com.nimbusds.jwt.JWTParser;
import org.springframework.boot.autoconfigure.condition.ConditionalOnProperty;
import org.springframework.security.authentication.BadCredentialsException;
import org.springframework.stereotype.Component;
import software.amazon.awssdk.services.cognitoidentityprovider.CognitoIdentityProviderClient;
import software.amazon.awssdk.services.cognitoidentityprovider.model.AuthFlowType;
import software.amazon.awssdk.services.cognitoidentityprovider.model.AuthenticationResultType;
import software.amazon.awssdk.services.cognitoidentityprovider.model.InitiateAuthRequest;
import software.amazon.awssdk.services.cognitoidentityprovider.model.InitiateAuthResponse;
import software.amazon.awssdk.services.cognitoidentityprovider.model.NotAuthorizedException;
import software.amazon.awssdk.services.cognitoidentityprovider.model.UserNotFoundException;

import java.text.ParseException;
import java.util.List;
import java.util.Map;

@Component
@ConditionalOnProperty(prefix = "app.cognito", name = "enabled", havingValue = "true")
public class CognitoUserAuthenticator {

    private final CognitoIdentityProviderClient cognitoClient;
    private final CognitoProperties cognitoProperties;

    public CognitoUserAuthenticator(CognitoIdentityProviderClient cognitoClient,
                                    CognitoProperties cognitoProperties) {
        this.cognitoClient = cognitoClient;
        this.cognitoProperties = cognitoProperties;
    }

    public AuthenticationResult authenticate(String username,
                                             String password,
                                             String clientId,
                                             String requiredRole) {
        try {
            InitiateAuthResponse response = cognitoClient.initiateAuth(InitiateAuthRequest.builder()
                    .authFlow(AuthFlowType.USER_PASSWORD_AUTH)
                    .clientId(clientId)
                    .authParameters(Map.of(
                            "USERNAME", username,
                            "PASSWORD", password
                    ))
                    .build());

            return toAuthenticationResult(response.authenticationResult(), username, requiredRole, null);
        } catch (NotAuthorizedException | UserNotFoundException ex) {
            throw new BadCredentialsException("Credenciales invalidas.");
        }
    }

    public AuthenticationResult refresh(String refreshToken, String clientId, String requiredRole) {
        try {
            InitiateAuthResponse response = cognitoClient.initiateAuth(InitiateAuthRequest.builder()
                    .authFlow(AuthFlowType.REFRESH_TOKEN_AUTH)
                    .clientId(clientId)
                    .authParameters(Map.of("REFRESH_TOKEN", refreshToken))
                    .build());

            return toAuthenticationResult(response.authenticationResult(), null, requiredRole, refreshToken);
        } catch (NotAuthorizedException ex) {
            throw new BadCredentialsException("Refresh token invalido o expirado.");
        }
    }

    private AuthenticationResult toAuthenticationResult(AuthenticationResultType authResult,
                                                          String fallbackUsername,
                                                          String requiredRole,
                                                          String fallbackRefreshToken) {
        if (authResult == null || authResult.accessToken() == null) {
            throw new BadCredentialsException("No se pudo autenticar con Cognito.");
        }

        List<String> groups = extractGroups(authResult.accessToken());
        List<String> roles = CognitoRoleMapper.toRoles(groups, cognitoProperties);
        if (roles.stream().noneMatch(requiredRole::equals)) {
            throw new BadCredentialsException("El usuario no tiene permisos para acceder.");
        }

        String username = extractUsername(authResult.accessToken(), fallbackUsername);
        String refreshToken = authResult.refreshToken() != null
                ? authResult.refreshToken()
                : fallbackRefreshToken;

        return new AuthenticationResult(authResult, username, roles, groups, refreshToken);
    }

    private String extractUsername(String accessToken, String fallback) {
        try {
            var claims = JWTParser.parse(accessToken).getJWTClaimsSet();
            String username = claims.getStringClaim("username");
            if (username == null || username.isBlank()) {
                username = claims.getStringClaim("cognito:username");
            }
            return username != null && !username.isBlank() ? username : fallback;
        } catch (ParseException ex) {
            throw new BadCredentialsException("No se pudo leer el usuario del token.");
        }
    }

    private List<String> extractGroups(String accessToken) {
        try {
            var claims = JWTParser.parse(accessToken).getJWTClaimsSet();
            List<String> groups = claims.getStringListClaim("cognito:groups");
            return groups != null ? groups : List.of();
        } catch (ParseException ex) {
            throw new BadCredentialsException("No se pudieron leer los grupos del token.");
        }
    }

    public record AuthenticationResult(
            AuthenticationResultType tokens,
            String username,
            List<String> roles,
            List<String> groups,
            String refreshToken
    ) {
    }
}