package com.yiboshi.science.config.security;


import com.yiboshi.science.service.SystemSetService;
import com.yiboshi.science.utils.SystemSetKey;
import lombok.AllArgsConstructor;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.security.core.AuthenticationException;
import org.springframework.security.oauth2.common.*;
import org.springframework.security.oauth2.provider.OAuth2Authentication;
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.transaction.annotation.Transactional;

import java.util.Date;
import java.util.UUID;

/**
 * @Auther: kylin
 * @Date: 2022/7/22 09:31
 * @Description:
 */
@AllArgsConstructor
public class AuthorizationTokenService extends DefaultTokenServices {

    private TokenEnhancer accessTokenEnhancer;

    @Autowired
    private TokenStore tokenStore;

    private SystemSetService systemSetService;

    public AuthorizationTokenService(SystemSetService systemSetService) {
        this.systemSetService = systemSetService;
    }

    @Transactional
    @Override
    public OAuth2AccessToken createAccessToken(OAuth2Authentication authentication) throws AuthenticationException {
        //  1允许同时登录一个账户 2 不允许同时登录有个账户
        String LoginType = systemSetService.getByKey(SystemSetKey.LoginType);
        if (LoginType.equals("1"))
            return super.createAccessToken(authentication);
        else {
            OAuth2AccessToken existingAccessToken = tokenStore.getAccessToken(authentication);
            OAuth2RefreshToken refreshToken = null;
            if (existingAccessToken != null) {
                if (existingAccessToken.getRefreshToken() != null) {
                    refreshToken = existingAccessToken.getRefreshToken();
                    tokenStore.removeRefreshToken(refreshToken);
                }
                tokenStore.removeAccessToken(existingAccessToken);
            }
            if (refreshToken == null) {
                refreshToken = createRefreshToken(authentication);
            } 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 value = UUID.randomUUID().toString();
        if (validitySeconds > 0) {
            return new DefaultExpiringOAuth2RefreshToken(value, new Date(System.currentTimeMillis()
                    + (validitySeconds * 1000L)));
        }
        return new DefaultOAuth2RefreshToken(value);
    }

    private OAuth2AccessToken createAccessToken(OAuth2Authentication authentication, OAuth2RefreshToken refreshToken) {
        DefaultOAuth2AccessToken token = new DefaultOAuth2AccessToken(UUID.randomUUID().toString());
        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 setTokenStore(TokenStore tokenStore) {
        super.setTokenStore(tokenStore);
        this.tokenStore = tokenStore;
    }

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