package org.radarcns.management.security;

import org.radarcns.management.domain.Source;
import org.radarcns.management.domain.Subject;
import org.radarcns.management.domain.User;
import org.radarcns.management.repository.SubjectRepository;
import org.radarcns.management.repository.UserRepository;
import org.springframework.beans.factory.InitializingBean;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.security.core.authority.AuthorityUtils;
import org.springframework.security.core.userdetails.UserDetails;
import org.springframework.security.oauth2.common.DefaultOAuth2AccessToken;
import org.springframework.security.oauth2.common.OAuth2AccessToken;
import org.springframework.security.oauth2.provider.OAuth2Authentication;
import org.springframework.security.oauth2.provider.token.TokenEnhancer;

import java.time.Instant;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Optional;
import java.util.Set;
import java.util.stream.Collectors;

public class ClaimsTokenEnhancer implements TokenEnhancer, InitializingBean {

    @Autowired
    private SubjectRepository subjectRepository;

    @Autowired
    private UserRepository userRepository;

    @Value("${spring.application.name}")
    private String appName;

    @Override
    public OAuth2AccessToken enhance(OAuth2AccessToken accessToken,
            OAuth2Authentication authentication) {
        if (authentication.getAuthorities() != null && !authentication.getAuthorities().isEmpty()) {
            Map<String, Object> additionalInfo = new HashMap<>();
            Set<String> currentUserAuthorities = AuthorityUtils.authorityListToSet(authentication.getAuthorities());

            String userName = null;
            if (authentication.getPrincipal() instanceof UserDetails) {
                UserDetails springSecurityUser = (UserDetails) authentication.getPrincipal();
                userName = springSecurityUser.getUsername();
            } else if (authentication.getPrincipal() instanceof String) {
                userName = (String) authentication.getPrincipal();
            }

            if(userName!=null) {
                // add the 'sub' claim in accordance with JWT spec
                additionalInfo.put("sub", userName);

                Optional<User> optUser = userRepository.findOneByLogin(userName);
                if (optUser.isPresent()) {
                    List<String> roles = optUser.get().getRoles().stream()
                        .filter(role -> role.getProject() != null)
                        .map(role -> role.getProject().getProjectName() + ":"
                            + role.getAuthority().getName())
                        .collect(Collectors.toList());
                    additionalInfo.put("roles", roles);

                }

                List<Source> assignedSources = subjectRepository
                    .findSourcesBySubjectLogin(userName);

                List<String> sourceIds = assignedSources.stream()
                    .map(s -> s.getSourceId().toString()).collect(Collectors.toList());
                additionalInfo.put("sources", sourceIds);
            }
            // add iat and iss optional JWT claims
            additionalInfo.put("iat", Instant.now().getEpochSecond());
            additionalInfo.put("iss", appName);
            ((DefaultOAuth2AccessToken) accessToken)
                .setAdditionalInformation(additionalInfo);
        }
        return accessToken;
    }

    @Override
    public void afterPropertiesSet() throws Exception {
    }
}

@@ -1492,297 +1492,72 @@
 ) %7B%0A
-        if (authentication.getAuthorities() != null && !authentication.getAuthorities().isEmpty()) %7B%0A            Map%3CString, Object%3E additionalInfo = new HashMap%3C%3E();%0A            Set%3CString%3E currentUserAuthorities = AuthorityUtils.authorityListToSet(authentication.getAuthorities());%0A%0A    
+%0A        Map%3CString, Object%3E additionalInfo = new HashMap%3C%3E();%0A%0A
     
@@ -1584,20 +1584,16 @@
 = null;%0A
-    
         
@@ -1664,20 +1664,16 @@
         
- 
-   
 UserDeta
@@ -1746,36 +1746,32 @@
 ();%0A            
-    
 userName = sprin
@@ -1795,28 +1795,24 @@
 Username();%0A
-    
         %7D el
@@ -1873,36 +1873,32 @@
 ) %7B%0A            
-    
 userName = (Stri
@@ -1940,34 +1940,27 @@
         
-    %7D%0A%0A   
+%7D%0A%0A
         
  if(user
@@ -1951,17 +1951,16 @@
         
- 
 if(userN
@@ -1968,28 +1968,24 @@
 me!=null) %7B%0A
-    
             
@@ -2039,36 +2039,32 @@
 pec%0A            
-    
 additionalInfo.p
@@ -2097,20 +2097,16 @@
         
-  
-  
 Optional
@@ -2159,28 +2159,24 @@
 (userName);%0A
-    
             
@@ -2210,36 +2210,32 @@
                 
-    
 List%3CString%3E rol
@@ -2293,20 +2293,16 @@
         
- 
-   
 .filter(
@@ -2336,20 +2336,16 @@
 = null)%0A
-    
         
@@ -2410,20 +2410,16 @@
 ) + %22:%22%0A
-    
         
@@ -2487,20 +2487,16 @@
         
-    
 .collect
@@ -2510,36 +2510,32 @@
 tors.toList());%0A
-    
                 
@@ -2579,35 +2579,27 @@
             
-   
- 
 %7D%0A%0A
-    
             
@@ -2663,20 +2663,16 @@
         
-    
 .findSou
@@ -2706,36 +2706,32 @@
 );%0A%0A            
-  
-  
 List%3CString%3E sou
@@ -2756,36 +2756,32 @@
 ources.stream()%0A
-    
                 
@@ -2852,36 +2852,32 @@
 ));%0A            
-    
 additionalInfo.p
@@ -2906,34 +2906,26 @@
 s);%0A        
- 
-   %7D%0A    
+%7D%0A
         // a
@@ -2959,20 +2959,16 @@
  claims%0A
-    
         
@@ -3026,36 +3026,32 @@
 ond());%0A        
-    
 additionalInfo.p
@@ -3066,28 +3066,24 @@
 , appName);%0A
-    
         ((De
@@ -3131,20 +3131,16 @@
         
-   
- 
 .setAddi
@@ -3174,25 +3174,16 @@
 lInfo);%0A
-        %7D
 %0A       

Enhance the JWT with all claims always