package org.apereo.cas.adaptors.yubikey.dao;

import org.apereo.cas.adaptors.yubikey.YubiKeyAccount;
import org.apereo.cas.adaptors.yubikey.YubiKeyAccountValidator;
import org.apereo.cas.adaptors.yubikey.YubiKeyDeviceRegistrationRequest;
import org.apereo.cas.adaptors.yubikey.YubiKeyRegisteredDevice;
import org.apereo.cas.adaptors.yubikey.registry.BaseYubiKeyAccountRegistry;

import lombok.val;
import org.springframework.data.mongodb.core.MongoOperations;
import org.springframework.data.mongodb.core.query.Criteria;
import org.springframework.data.mongodb.core.query.Query;
import org.springframework.data.mongodb.core.query.Update;

import java.util.Arrays;
import java.util.Collection;
import java.util.stream.Collectors;

/**
 * This is {@link MongoDbYubiKeyAccountRegistry}.
 *
 * @author Misagh Moayyed
 * @since 5.2.0
 */
public class MongoDbYubiKeyAccountRegistry extends BaseYubiKeyAccountRegistry {

    private final String collectionName;

    private final MongoOperations mongoTemplate;

    public MongoDbYubiKeyAccountRegistry(final YubiKeyAccountValidator accountValidator,
                                         final MongoOperations mongoTemplate,
                                         final String collectionName) {
        super(accountValidator);
        this.mongoTemplate = mongoTemplate;
        this.collectionName = collectionName;
    }

    @Override
    public YubiKeyAccount getAccountInternal(final String uid) {
        val query = new Query().addCriteria(Criteria.where(MongoDbYubiKeyAccount.FIELD_USERNAME).is(uid));
        return mongoTemplate.findOne(query, MongoDbYubiKeyAccount.class, this.collectionName);
    }

    @Override
    public void delete(final String uid) {
        val query = new Query();
        query.addCriteria(Criteria.where(MongoDbYubiKeyAccount.FIELD_USERNAME).is(uid));
        this.mongoTemplate.remove(query, MongoDbYubiKeyAccount.class, this.collectionName);
    }

    @Override
    public void delete(final String username, final long deviceId) {
        val query = new Query();
        query.addCriteria(
            Criteria.where(MongoDbYubiKeyAccount.FIELD_USERNAME).is(username)
                .and(MongoDbYubiKeyAccount.FIELD_DEVICES).elemMatch(Criteria.where("_id").is(deviceId)));
        this.mongoTemplate.remove(query, this.collectionName);
    }

    @Override
    public void deleteAll() {
        this.mongoTemplate.remove(new Query(), MongoDbYubiKeyAccount.class, this.collectionName);
    }

    @Override
    protected Collection<? extends YubiKeyAccount> getAccountsInternal() {
        return this.mongoTemplate.findAll(MongoDbYubiKeyAccount.class, this.collectionName);
    }

    @Override
    public boolean update(final YubiKeyAccount account) {
        val query = new Query().addCriteria(Criteria.where(MongoDbYubiKeyAccount.FIELD_USERNAME).is(account.getUsername()));
        val update = Update.update(MongoDbYubiKeyAccount.FIELD_DEVICES, account.getDevices());
        this.mongoTemplate.updateFirst(query, update, MongoDbYubiKeyAccount.class, this.collectionName);
        return true;
    }

    @Override
    public YubiKeyAccount save(final YubiKeyDeviceRegistrationRequest request, final YubiKeyRegisteredDevice... device) {
        val result = MongoDbYubiKeyAccount.builder()
            .id(System.currentTimeMillis())
            .username(request.getUsername())
            .devices(Arrays.stream(device).collect(Collectors.toList()))
            .build();
        return mongoTemplate.save(result, this.collectionName);
    }
}

@@ -2033,123 +2033,36 @@
 val 
-query = new Query();%0A        query.addCriteria(%0A            Criteria.where(MongoDbYubiKeyAccount.FIELD_USERNAME).is
+account = getAccountInternal
 (use
@@ -2067,29 +2067,26 @@
 sername)
+;
 %0A
-    
         
     .and
@@ -2081,89 +2081,70 @@
     
-    .and(MongoDbYubiKeyAccount.FIELD_DEVICES).elemMatch(Criteria.where(%22_id%22).is(
+if (account.getDevices().removeIf(device -%3E device.getId() == 
 devi
@@ -2149,27 +2149,28 @@
 viceId))
-);%0A
+ %7B%0A 
         
 this.mon
@@ -2165,62 +2165,37 @@
     
-this.mongoTemplate.remove(query, this.collectionName);
+   update(account);%0A        %7D
 %0A   
@@ -2351,197 +2351,8 @@
  %7D%0A%0A
-    @Override%0A    protected Collection%3C? extends YubiKeyAccount%3E getAccountsInternal() %7B%0A        return this.mongoTemplate.findAll(MongoDbYubiKeyAccount.class, this.collectionName);%0A    %7D%0A%0A
     
@@ -3214,11 +3214,200 @@
 );%0A    %7D
+%0A%0A    @Override%0A    protected Collection%3C? extends YubiKeyAccount%3E getAccountsInternal() %7B%0A        return this.mongoTemplate.findAll(MongoDbYubiKeyAccount.class, this.collectionName);%0A    %7D
 %0A%7D%0A

fix mongo yubikey registry