package com.keyware.shandan.config;

import com.keyware.shandan.common.util.StringUtils;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.security.core.Authentication;
import org.springframework.security.oauth2.common.OAuth2AccessToken;
import org.springframework.security.oauth2.provider.OAuth2Authentication;
import org.springframework.security.web.authentication.logout.LogoutHandler;
import org.springframework.security.web.authentication.logout.SecurityContextLogoutHandler;
import org.springframework.stereotype.Component;

import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletResponse;

@Component
public class SecurityLogoutHandler implements LogoutHandler {

    @Autowired
    private InMemoryTokenStoreImpl inMemoryTokenStore;

    private final SecurityContextLogoutHandler logoutHandler = new SecurityContextLogoutHandler();

    @Override
    public void logout(HttpServletRequest request, HttpServletResponse response, Authentication authentication) {
        String author = request.getHeader("Authorization");
        String loginName = request.getParameter("loginName");
        if (authentication != null) {
            logoutHandler.logout(request, response, authentication);
        } else if (StringUtils.isNotBlank(author)) {
            author = author.replaceFirst("Bearer ", "");
            OAuth2AccessToken token = inMemoryTokenStore.readAccessToken(author);
            if (token != null) {
                logout(request, response, token);
            }
        } else if (StringUtils.hasText(loginName)) {
            inMemoryTokenStore.findTokensByClientIdAndUserName("desktop", loginName).forEach(token -> logout(request, response, token));
        }
    }

    private void logout(HttpServletRequest request, HttpServletResponse response, OAuth2AccessToken token) {
        inMemoryTokenStore.removeAccessToken(token);
        inMemoryTokenStore.removeRefreshToken(token.getRefreshToken());
        inMemoryTokenStore.removeAccessTokenUsingRefreshToken(token.getRefreshToken());
        OAuth2Authentication oauth2Auth = inMemoryTokenStore.readAuthentication(token);
        if (oauth2Auth != null) {
            Authentication auth = oauth2Auth.getUserAuthentication();
            logoutHandler.logout(request, response, auth);
        }
    }
}
