package ee.tuleva.onboarding.auth;

import com.codeborne.security.mobileid.MobileIDSession;
import ee.tuleva.onboarding.user.User;
import ee.tuleva.onboarding.user.UserRepository;
import lombok.extern.slf4j.Slf4j;
import org.springframework.security.core.Authentication;
import org.springframework.security.oauth2.common.OAuth2AccessToken;
import org.springframework.security.oauth2.common.exceptions.InvalidRequestException;
import org.springframework.security.oauth2.provider.*;
import org.springframework.security.oauth2.provider.token.AbstractTokenGranter;
import org.springframework.security.oauth2.provider.token.AuthorizationServerTokenServices;

@Slf4j
public class MobileIdTokenGranter extends AbstractTokenGranter implements TokenGranter {
    public static final String GRANT_TYPE = "mobile_id";

    private final MobileIdAuthService mobileIdAuthService;

    private final UserRepository userRepository;

    public MobileIdTokenGranter(AuthorizationServerTokenServices tokenServices,
                                ClientDetailsService clientDetailsService,
                                OAuth2RequestFactory requestFactory,
                                MobileIdAuthService mobileIdAuthService,
                                UserRepository userRepository) {
        super(tokenServices, clientDetailsService, requestFactory, GRANT_TYPE);
        this.mobileIdAuthService = mobileIdAuthService;
        this.userRepository = userRepository;
    }

    @Override
    protected OAuth2AccessToken getAccessToken(ClientDetails client, TokenRequest tokenRequest) {
        // grant_type validated in AbstractTokenGranter

        assert mobileIdAuthService != null;

        final String clientId = client.getClientId();
        if (clientId == null) {
            throw new InvalidRequestException("Unknown Client ID.");
        }

        MobileIDSession mobileIDSession = MobileIdSessionStore.get();
        boolean isComplete = mobileIdAuthService.isLoginComplete(mobileIDSession);

        if(!isComplete) {
            throw new MobileIdAuthNotCompleteException();
        }

        User user = userRepository.findByPersonalCode(mobileIDSession.personalCode);
        Authentication userAuthentication = new PersonalCodeAuthentication(user, mobileIDSession, null);
        userAuthentication.setAuthenticated(true);

/*

        if (userAuthentication == null) {
            log.error("Failed to authenticate user: got null authentication");
            throw new InvalidGrantException("Invalid user credentials.");
        }
*/

        final OAuth2Request oAuth2Request = tokenRequest.createOAuth2Request(client);
        final OAuth2Authentication oAuth2Authentication = new OAuth2Authentication(oAuth2Request,
                userAuthentication
        );

        return getTokenServices().createAccessToken(oAuth2Authentication);
    }
}

@@ -1754,32 +1754,101 @@
 ntId == null) %7B%0A
+            log.error(%22Failed to authenticate client %7B%7D%22, clientId);%0A
             thro
@@ -2245,24 +2245,25 @@
 lCode);%0A
+%0A
         
 Authenti
@@ -2258,173 +2258,264 @@
     
-Authentication userAuthentication = new P
+if (user == null) %7B%0A            log.error(%22Failed to authenticate user: couldn't find user with p
 ersonal
-CodeAuthentication(user, mobileIDSession, null);%0A        userAuthentication.setAuthenticated(true);%0A%0A/*%0A%0A        if (
+ code %7B%7D%22, mobileIDSession.personalCode);%0A            throw new InvalidRequestException(%22Invalid user credentials.%22);%0A        %7D%0A%0A        Authentication 
 user
@@ -2534,183 +2534,120 @@
 on =
-=
  n
-ull) %7B%0A            log.error(%22Failed to authenticate user: got null authentication%22);%0A            throw new InvalidGrantException(%22Invalid user credentials.%22);%0A        %7D%0A*/
+ew PersonalCodeAuthentication(user, mobileIDSession, null);%0A        userAuthentication.setAuthenticated(true);
 %0A%0A  

Add exception handling for unknown users