package org.opencb.opencga.catalog.db.mongodb;

import com.mongodb.client.model.Filters;
import org.bson.Document;
import org.bson.conversions.Bson;
import org.opencb.commons.datastore.core.DataResult;
import org.opencb.commons.datastore.core.Query;
import org.opencb.commons.datastore.core.QueryOptions;
import org.opencb.commons.datastore.mongodb.MongoDBCollection;
import org.opencb.opencga.catalog.db.api.MigrationDBAdaptor;
import org.opencb.opencga.catalog.db.mongodb.converters.MigrationConverter;
import org.opencb.opencga.catalog.exceptions.CatalogDBException;
import org.opencb.opencga.catalog.migration.MigrationRun;
import org.opencb.opencga.core.config.Configuration;
import org.opencb.opencga.core.response.OpenCGAResult;
import org.slf4j.LoggerFactory;

import java.util.ArrayList;
import java.util.List;
import java.util.Map;

public class MigrationMongoDBAdaptor extends MongoDBAdaptor implements MigrationDBAdaptor {

    private final MongoDBCollection migrationCollection;
    private final MigrationConverter migrationConverter;

    public MigrationMongoDBAdaptor(MongoDBCollection migrationCollection, Configuration configuration,
                                   MongoDBAdaptorFactory dbAdaptorFactory) {
        super(configuration, LoggerFactory.getLogger(MigrationMongoDBAdaptor.class));
        this.dbAdaptorFactory = dbAdaptorFactory;
        this.migrationCollection = migrationCollection;
        this.migrationConverter = new MigrationConverter();
    }

    @Override
    public void upsert(MigrationRun migrationRun) throws CatalogDBException {
        Document migrationDocument = migrationConverter.convertToStorageType(migrationRun);
        if (get(migrationRun.getId()).getNumResults() == 0) {
            DataResult insert = migrationCollection.insert(migrationDocument, QueryOptions.empty());
            if (insert.getNumUpdated() == 0) {
                throw new CatalogDBException("Could not insert MigrationRun");
            }
        } else {
            Query query = new Query()
                    .append(QueryParams.ID.key(), migrationRun.getId())
                    .append(QueryParams.VERSION.key(), migrationRun.getVersion());
            Bson bsonQuery = parseQuery(query);
            DataResult update = migrationCollection.update(bsonQuery, new Document("$set", migrationDocument), QueryOptions.empty());
            if (update.getNumUpdated() == 0) {
                throw new CatalogDBException("Could not update MigrationRun");
            }
        }
    }

    @Override
    public OpenCGAResult<MigrationRun> get(Query query) throws CatalogDBException {
        Bson bsonQuery = parseQuery(query);
        return new OpenCGAResult<>(migrationCollection.find(bsonQuery, migrationConverter, QueryOptions.empty()));
    }

    @Override
    public OpenCGAResult<MigrationRun> get(List<String> migrationRunIds) throws CatalogDBException {
        Query query = new Query(QueryParams.ID.key(), migrationRunIds);
        return get(query);
    }

    private Bson parseQuery(Query query) throws CatalogDBException {
        List<Bson> andBsonList = new ArrayList<>();
        Query queryCopy = new Query(query);

        for (Map.Entry<String, Object> entry : queryCopy.entrySet()) {
            String key = entry.getKey().split("\\.")[0];
            MigrationDBAdaptor.QueryParams queryParam = MigrationDBAdaptor.QueryParams.getParam(entry.getKey()) != null
                    ? MigrationDBAdaptor.QueryParams.getParam(entry.getKey())
                    : MigrationDBAdaptor.QueryParams.getParam(key);
            if (queryParam == null) {
                throw new CatalogDBException("Unexpected parameter " + entry.getKey() + ". The parameter does not exist or cannot be "
                        + "queried for.");
            }
            try {
                switch (queryParam) {
                    case ID:
                    case VERSION:
                    case START:
                    case END:
                    case PATCH:
                    case STATUS:
                        addAutoOrQuery(queryParam.key(), queryParam.key(), queryCopy, queryParam.type(), andBsonList);
                        break;
                    default:
                        throw new CatalogDBException("Cannot query by parameter " + queryParam.key());
                }
            } catch (Exception e) {
                throw new CatalogDBException(e);
            }
        }

        if (andBsonList.size() > 0) {
            return Filters.and(andBsonList);
        } else {
            return new Document();
        }
    }

}

@@ -1671,332 +1671,8 @@
 n);%0A
-        if (get(migrationRun.getId()).getNumResults() == 0) %7B%0A            DataResult insert = migrationCollection.insert(migrationDocument, QueryOptions.empty());%0A            if (insert.getNumUpdated() == 0) %7B%0A                throw new CatalogDBException(%22Could not insert MigrationRun%22);%0A            %7D%0A        %7D else %7B%0A    
     
@@ -1701,20 +1701,16 @@
 Query()%0A
-    
         
@@ -1777,36 +1777,32 @@
                 
-    
 .append(QueryPar
@@ -1852,28 +1852,24 @@
 ));%0A        
-   
- 
 Bson bsonQue
@@ -1888,28 +1888,24 @@
 ery(query);%0A
-    
         Data
@@ -2002,36 +2002,76 @@
 nt),
- QueryOptions.empty());%0A    
+%0A                new QueryOptions(MongoDBCollection.UPSERT, true));%0A
     
@@ -2105,24 +2105,52 @@
 d() == 0
+ && update.getNumInserted() == 0
 ) %7B%0A
-    
         
@@ -2212,38 +2212,24 @@
 ationRun%22);%0A
-            %7D%0A
         %7D%0A  

Use actual upsert for updating migration runs