package cc.mrbird.febs.auth.configure;

import com.ruoyi.common.core.utils.ServletUtils;
import org.apache.commons.codec.binary.Base64;
import org.apache.commons.lang3.StringUtils;
import org.springframework.security.authentication.AuthenticationManager;
import org.springframework.security.core.Authentication;
import org.springframework.security.core.AuthenticationException;
import org.springframework.security.crypto.keygen.BytesKeyGenerator;
import org.springframework.security.crypto.keygen.KeyGenerators;
import org.springframework.security.oauth2.common.*;
import org.springframework.security.oauth2.common.exceptions.InvalidGrantException;
import org.springframework.security.oauth2.common.exceptions.InvalidScopeException;
import org.springframework.security.oauth2.common.exceptions.InvalidTokenException;
import org.springframework.security.oauth2.provider.OAuth2Authentication;
import org.springframework.security.oauth2.provider.OAuth2Request;
import org.springframework.security.oauth2.provider.TokenRequest;
import org.springframework.security.oauth2.provider.token.DefaultTokenServices;
import org.springframework.security.oauth2.provider.token.TokenEnhancer;
import org.springframework.security.oauth2.provider.token.TokenStore;
import org.springframework.security.web.authentication.preauth.PreAuthenticatedAuthenticationToken;
import org.springframework.transaction.annotation.Transactional;

import java.nio.charset.Charset;
import java.util.Date;
import java.util.Set;

/**
 * 多模块生成的token加前缀
 */
public class CustomTokenServices extends DefaultTokenServices {

    private static final BytesKeyGenerator DEFAULT_TOKEN_GENERATOR = KeyGenerators.secureRandom(20);

    private static final Charset US_ASCII = Charset.forName("US-ASCII");

    private TokenStore tokenStore;
    private TokenEnhancer accessTokenEnhancer;

    private boolean supportRefreshToken = false;

    private boolean reuseRefreshToken = true;

    private AuthenticationManager authenticationManager;

    // 新增TransportClient依赖
    private TransportClient transportClient;

    // 新增TransportClient的setter方法
    public void setTransportClient(TransportClient transportClient) {
        this.transportClient = transportClient;
    }

    @Override
    public OAuth2AccessToken createAccessToken(OAuth2Authentication authentication) throws AuthenticationException {
        OAuth2AccessToken existingAccessToken = tokenStore.getAccessToken(authentication);
        OAuth2RefreshToken refreshToken = null;
        if (existingAccessToken != null) {
            if (existingAccessToken.isExpired()) {
                if (existingAccessToken.getRefreshToken() != null) {
                    refreshToken = existingAccessToken.getRefreshToken();
                    // The token store could remove the refresh token when the
                    // access token is removed, but we want to
                    // be sure...
                    tokenStore.removeRefreshToken(refreshToken);
                }
                tokenStore.removeAccessToken(existingAccessToken);
            } else {
                // Re-store the access token in case the authentication has changed
                tokenStore.storeAccessToken(existingAccessToken, authentication);
                return existingAccessToken;
            }
        }

        // Only create a new refresh token if there wasn't an existing one
        // associated with an expired access token.
        // Clients might be holding existing refresh tokens, so we re-use it in
        // the case that the old access token
        // expired.
        if (refreshToken == null) {
            refreshToken = createRefreshToken(authentication);
        }
        // But the refresh token itself might need to be re-issued if it has
        // expired.
        else if (refreshToken instanceof ExpiringOAuth2RefreshToken) {
            ExpiringOAuth2RefreshToken expiring = (ExpiringOAuth2RefreshToken) refreshToken;
            if (System.currentTimeMillis() > expiring.getExpiration().getTime()) {
                refreshToken = createRefreshToken(authentication);
            }
        }

        OAuth2AccessToken accessToken = createAccessToken(authentication, refreshToken);
        tokenStore.storeAccessToken(accessToken, authentication);
        // In case it was modified
        refreshToken = accessToken.getRefreshToken();
        if (refreshToken != null) {
            tokenStore.storeRefreshToken(refreshToken, authentication);
        }
        return accessToken;
    }

    private OAuth2RefreshToken createRefreshToken(OAuth2Authentication authentication) {
        if (!isSupportRefreshToken(authentication.getOAuth2Request())) {
            return null;
        }
        int validitySeconds = getRefreshTokenValiditySeconds(authentication.getOAuth2Request());

        String tokenValue = "";
        // 修改点1：根据认证模式区分来源
        String grantType = authentication.getOAuth2Request().getGrantType();
        boolean isClientCredential = "client_credentials".equals(grantType);

        if (isClientCredential) {
            Object principal = authentication.getPrincipal();
            boolean isClientPrincipal = principal != null && principal.toString().toLowerCase().contains("client");
            if (isClientPrincipal) {
                tokenValue = transportClient.calculateSM3(new String(Base64.encodeBase64URLSafe(DEFAULT_TOKEN_GENERATOR.generateKey()), US_ASCII));
            } else {

                try {
                    tokenValue = transportClient.getToken(); // 假设TransportClient有getToken()方法
                } catch (Exception e) {
                    e.printStackTrace();
                    throw new AuthenticationException("Failed to generate token via TransportClient", e) {
                    };
                }
            }
        } else {
            tokenValue = new String(Base64.encodeBase64URLSafe(DEFAULT_TOKEN_GENERATOR.generateKey()), US_ASCII);
        }

        if (validitySeconds > 0) {
            return new DefaultExpiringOAuth2RefreshToken(tokenValue, new Date(System.currentTimeMillis()
                    + (validitySeconds * 1000L)));
        }
        return new DefaultOAuth2RefreshToken(tokenValue);
    }

    private OAuth2AccessToken createAccessToken(OAuth2Authentication authentication, OAuth2RefreshToken refreshToken) {

        String sourceType = "";
        String tokenValue = "";
        // 修改点1：根据认证模式区分来源
        String grantType = authentication.getOAuth2Request().getGrantType();
        boolean isClientCredential = "client_credentials".equals(grantType);

//        String sourceType = ServletUtils.getRequest().getHeader("SourceType");
//        // 第三方客户端调用
//        if(StringUtils.isEmpty(sourceType)) {
//            sourceType = "data_exchange";
//        }
        //密码模式
//        if(authentication.getUserAuthentication() != null && authentication.getUserAuthentication() instanceof UsernamePasswordAuthenticationToken) {
//            UsernamePasswordAuthenticationToken userAuthentication = (UsernamePasswordAuthenticationToken) authentication.getUserAuthentication();
        // 读取请求数据
//            LinkedHashMap details = (LinkedHashMap) userAuthentication.getDetails();
//            sourceType = details.get("sourceType") + "_";
//        }
        if (isClientCredential) {
            Object principal = authentication.getPrincipal();
            boolean isClientPrincipal = principal != null && principal.toString().toLowerCase().contains("client");
            if (isClientPrincipal) {
                sourceType = "data_exchange_client";
                tokenValue = new String(Base64.encodeBase64URLSafe(DEFAULT_TOKEN_GENERATOR.generateKey()), US_ASCII);
            } else {
                sourceType = "data_exchange_feign";
                try {
                    tokenValue = transportClient.getToken(); // 假设TransportClient有getToken()方法
                } catch (Exception e) {
                    e.printStackTrace();
                    throw new AuthenticationException("Failed to generate token via TransportClient", e) {
                    };
                }
            }
        } else {
            sourceType = "data_exchange_login";
            tokenValue = new String(Base64.encodeBase64URLSafe(DEFAULT_TOKEN_GENERATOR.generateKey()), US_ASCII);
        }
        DefaultOAuth2AccessToken token = new DefaultOAuth2AccessToken(sourceType + "_" + tokenValue);
        if ("data_exchange_client".equals(sourceType)) {
            token = new DefaultOAuth2AccessToken(transportClient.calculateSM3(sourceType + "_" + tokenValue));
        }

        int validitySeconds = getAccessTokenValiditySeconds(authentication.getOAuth2Request());
        if (validitySeconds > 0) {
            token.setExpiration(new Date(System.currentTimeMillis() + (validitySeconds * 1000L)));
        }
        token.setRefreshToken(refreshToken);
        token.setScope(authentication.getOAuth2Request().getScope());

        return accessTokenEnhancer != null ? accessTokenEnhancer.enhance(token, authentication) : token;
    }

    @Override
    public void setTokenEnhancer(TokenEnhancer accessTokenEnhancer) {
        super.setTokenEnhancer(accessTokenEnhancer);
        this.accessTokenEnhancer = accessTokenEnhancer;
    }

    @Override
    public void setTokenStore(TokenStore tokenStore) {
        super.setTokenStore(tokenStore);
        this.tokenStore = tokenStore;
    }

    @Override
    @Transactional(noRollbackFor = {InvalidTokenException.class, InvalidGrantException.class})
    public OAuth2AccessToken refreshAccessToken(String refreshTokenValue, TokenRequest tokenRequest)
            throws AuthenticationException {

        if (!supportRefreshToken) {
            throw new InvalidGrantException("Invalid refresh token");
        }

        OAuth2RefreshToken refreshToken = tokenStore.readRefreshToken(refreshTokenValue);
        if (refreshToken == null) {
            throw new InvalidGrantException("Invalid refresh token");
        }

        OAuth2Authentication authentication = tokenStore.readAuthenticationForRefreshToken(refreshToken);
        if (this.authenticationManager != null && !authentication.isClientOnly()) {
            // The client has already been authenticated, but the user authentication might be old now, so give it a
            // chance to re-authenticate.
            Authentication userAuthentication = authentication.getUserAuthentication();
            PreAuthenticatedAuthenticationToken preAuthenticatedToken = new PreAuthenticatedAuthenticationToken(
                    userAuthentication,
                    "",
                    authentication.getAuthorities()
            );
            if (userAuthentication.getDetails() != null) {
                preAuthenticatedToken.setDetails(userAuthentication.getDetails());
            }
            Authentication user = authenticationManager.authenticate(preAuthenticatedToken);
            Object details = authentication.getDetails();
            authentication = new OAuth2Authentication(authentication.getOAuth2Request(), user);
            authentication.setDetails(details);
        }
        String clientId = authentication.getOAuth2Request().getClientId();
        if (clientId == null || !clientId.equals(tokenRequest.getClientId())) {
            throw new InvalidGrantException("Wrong client for this refresh token");
        }

        String orgAccessToken = tokenStore.getAccessToken(authentication).getValue();

        // clear out any access tokens already associated with the refresh
        // token.
        tokenStore.removeAccessTokenUsingRefreshToken(refreshToken);

        if (isExpired(refreshToken)) {
            tokenStore.removeRefreshToken(refreshToken);
            throw new InvalidTokenException("Invalid refresh token (expired)");
        }

        authentication = createRefreshedAuthentication(authentication, tokenRequest);

        if (!reuseRefreshToken) {
            tokenStore.removeRefreshToken(refreshToken);
            refreshToken = createRefreshToken(authentication);
        }

        OAuth2AccessToken accessToken = createAccessToken(authentication, refreshToken);
        // 保留token原值
        ((DefaultOAuth2AccessToken) accessToken).setValue(orgAccessToken);
        tokenStore.storeAccessToken(accessToken, authentication);
        if (!reuseRefreshToken) {
            tokenStore.storeRefreshToken(accessToken.getRefreshToken(), authentication);
        }
        return accessToken;
    }

    private OAuth2Authentication createRefreshedAuthentication(OAuth2Authentication authentication, TokenRequest request) {
        OAuth2Authentication narrowed = authentication;
        Set<String> scope = request.getScope();
        OAuth2Request clientAuth = authentication.getOAuth2Request().refresh(request);
        if (scope != null && !scope.isEmpty()) {
            Set<String> originalScope = clientAuth.getScope();
            if (originalScope == null || !originalScope.containsAll(scope)) {
                throw new InvalidScopeException("Unable to narrow the scope of the client authentication", originalScope);
            } else {
                clientAuth = clientAuth.narrowScope(scope);
            }
        }
        narrowed = new OAuth2Authentication(clientAuth, authentication.getUserAuthentication());
        return narrowed;
    }

    /**
     * Whether to support the refresh token.
     *
     * @param supportRefreshToken Whether to support the refresh token.
     */
    public void setSupportRefreshToken(boolean supportRefreshToken) {
        this.supportRefreshToken = supportRefreshToken;
    }
}
