/*
 * Copyright 1999-2021 Alibaba Group Holding Ltd.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *      http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package com.alibaba.nacos.plugin.auth.impl.token.impl;

import com.alibaba.nacos.plugin.auth.exception.AccessException;
import com.alibaba.nacos.plugin.auth.impl.token.TokenManager;
import com.alibaba.nacos.plugin.auth.impl.users.NacosUser;
import org.springframework.scheduling.annotation.Scheduled;
import org.springframework.security.core.Authentication;

import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.TimeUnit;

/**
 * Cached JWT token manager.
 *
 * @author majorhe
 */
public class CachedJwtTokenManager implements TokenManager {
    
    /**
     * key: token string, value: token entity.
     */
    private volatile Map<String, TokenEntity> tokenMap = new ConcurrentHashMap<>(1024);
    
    /**
     * key: username, value: token entity. cache token created by self.
     */
    private volatile Map<String, TokenEntity> userMap = new ConcurrentHashMap<>(128);
    
    private final JwtTokenManager jwtTokenManager;
    
    public CachedJwtTokenManager(JwtTokenManager jwtTokenManager) {
        this.jwtTokenManager = jwtTokenManager;
    }
    
    @Scheduled(initialDelay = 30000, fixedDelay = 60000)
    private void cleanExpiredToken() {
        List<String> tokens = new ArrayList<>();
        tokenMap.forEach((k, v) -> {
            if (v.getExpiredTimeMills() < System.currentTimeMillis()) {
                tokens.add(k);
            }
        });
        tokens.forEach(e -> tokenMap.remove(e));
        List<String> users = new ArrayList<>();
        userMap.forEach((k, v) -> {
            if (v.getExpiredTimeMills() < System.currentTimeMillis()) {
                users.add(k);
            }
        });
        users.forEach(e -> userMap.remove(e));
    }
    
    @Override
    public String createToken(Authentication authentication) throws AccessException {
        return createToken(authentication.getName());
    }
    
    /**
     * Create token.
     *
     * @param username auth info
     * @return token
     * @throws AccessException access exception
     */
    public String createToken(String username) throws AccessException {
        if (userMap.containsKey(username)) {
            String token = userMap.get(username).getToken();
            long expiredTime = userMap.get(username).getExpiredTimeMills();
            if (!needRefresh(expiredTime)) {
                return token;
            }
        }
        String token = jwtTokenManager.createToken(username);
        NacosUser user = jwtTokenManager.parseToken(token);
        long expiredTime = System.currentTimeMillis() + TimeUnit.SECONDS.toMillis(getTokenValidityInSeconds());
        Authentication authentication = jwtTokenManager.getAuthentication(token);
        TokenEntity model = new TokenEntity(token, username, expiredTime, authentication, user);
        tokenMap.put(token, model);
        userMap.put(username, model);
        return token;
    }
    
    /**
     * Get auth Info.
     *
     * @param token token
     * @return auth info
     * @throws AccessException access exception
     */
    public Authentication getAuthentication(String token) throws AccessException {
        if (!tokenMap.containsKey(token)) {
            return jwtTokenManager.getAuthentication(token);
        }
        return tokenMap.get(token).getAuthentication();
    }
    
    /**
     * validate token.
     *
     * @param token token
     * @throws AccessException access exception
     */
    public void validateToken(String token) throws AccessException {
        if (!tokenMap.containsKey(token)) {
            // jwtTokenManager.validateToken(token) will throw runtime exception if token invalid
            jwtTokenManager.validateToken(token);
            // if token valid
            Authentication authentication = jwtTokenManager.getAuthentication(token);
            String username = authentication.getName();
            if (username == null || username.isEmpty()) {
                return;
            }
            long expiredTime = TimeUnit.SECONDS.toMillis(jwtTokenManager.getExpiredTimeInSeconds(token));
            if (expiredTime <= System.currentTimeMillis()) {
                return;
            }
            NacosUser user = jwtTokenManager.parseToken(token);
            tokenMap.putIfAbsent(token, new TokenEntity(token, username, expiredTime, authentication, user));
        }
    }
    
    @Override
    public NacosUser parseToken(String token) throws AccessException {
        if (!tokenMap.containsKey(token)) {
            Authentication authentication = jwtTokenManager.getAuthentication(token);
            String username = authentication.getName();
            if (username == null || username.isEmpty()) {
                throw new AccessException("invalid token, username is empty");
            }
            long expiredTime = TimeUnit.SECONDS.toMillis(jwtTokenManager.getExpiredTimeInSeconds(token));
            if (expiredTime <= System.currentTimeMillis()) {
                throw new AccessException("expired token");
            }
            NacosUser user = jwtTokenManager.parseToken(token);
            tokenMap.putIfAbsent(token, new TokenEntity(token, username, expiredTime, authentication, user));
            return user;
        }
        return tokenMap.get(token).getNacosUser();
    }
    
    public long getTokenTtlInSeconds(String token) throws AccessException {
        if (tokenMap.containsKey(token)) {
            return TimeUnit.MILLISECONDS.toSeconds(
                    tokenMap.get(token).getExpiredTimeMills() - System.currentTimeMillis());
        }
        return jwtTokenManager.getTokenTtlInSeconds(token);
    }
    
    @Override
    public long getTokenValidityInSeconds() {
        return jwtTokenManager.getTokenValidityInSeconds();
    }
    
    private boolean needRefresh(long expiredTimeMills) {
        long refreshWindowMills = TimeUnit.SECONDS.toMillis(getTokenValidityInSeconds() / 10);
        return System.currentTimeMillis() + refreshWindowMills > expiredTimeMills;
    }
    
    static class TokenEntity {
        
        private String token;
        
        private String userName;
        
        private long expiredTimeMills;
        
        private Authentication authentication;
        
        private NacosUser nacosUser;
        
        public TokenEntity(String token, String userName, long expiredTimeMills, Authentication authentication,
                NacosUser nacosUser) {
            this.token = token;
            this.userName = userName;
            this.expiredTimeMills = expiredTimeMills;
            this.authentication = authentication;
            this.nacosUser = nacosUser;
        }
        
        public String getToken() {
            return token;
        }
        
        public void setToken(String token) {
            this.token = token;
        }
        
        public String getUserName() {
            return userName;
        }
        
        public void setUserName(String userName) {
            this.userName = userName;
        }
        
        public long getExpiredTimeMills() {
            return expiredTimeMills;
        }
        
        public void setExpiredTimeMills(long expiredTimeMills) {
            this.expiredTimeMills = expiredTimeMills;
        }
        
        public Authentication getAuthentication() {
            return authentication;
        }
        
        public void setAuthentication(Authentication authentication) {
            this.authentication = authentication;
        }
        
        public NacosUser getNacosUser() {
            return nacosUser;
        }
        
        public void setNacosUser(NacosUser nacosUser) {
            this.nacosUser = nacosUser;
        }
        
        @Override
        public String toString() {
            return "TokenEntity{" + "token='" + token + '\'' + ", userName='" + userName + '\'' + ", expiredTimeMills="
                    + expiredTimeMills + ", authentication=" + authentication + ", nacosUser=" + nacosUser + '}';
        }
    }
    
}
