package org.opencb.opencga.app.migrations.v2_4_2.catalog;

import com.mongodb.client.MongoCollection;
import com.mongodb.client.MongoCursor;
import com.mongodb.client.model.Filters;
import com.mongodb.client.model.Projections;
import com.mongodb.client.model.Sorts;
import org.apache.commons.collections4.CollectionUtils;
import org.bson.Document;
import org.opencb.commons.datastore.core.ObjectMap;
import org.opencb.commons.datastore.core.QueryOptions;
import org.opencb.opencga.catalog.db.api.ClinicalAnalysisDBAdaptor;
import org.opencb.opencga.catalog.db.api.IndividualDBAdaptor;
import org.opencb.opencga.catalog.db.mongodb.MongoDBAdaptorFactory;
import org.opencb.opencga.catalog.db.mongodb.converters.ClinicalAnalysisConverter;
import org.opencb.opencga.catalog.exceptions.CatalogException;
import org.opencb.opencga.catalog.migration.Migration;
import org.opencb.opencga.catalog.migration.MigrationTool;
import org.opencb.opencga.core.models.clinical.ClinicalAnalysis;
import org.opencb.opencga.core.models.common.Enums;
import org.opencb.opencga.core.models.individual.Individual;
import org.opencb.opencga.core.models.sample.Sample;

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

@Migration(id = "recover_proband_samples_in_cases_TASK-1470",
        description = "Recover lost samples in clinical collection #TASK-1470", version = "2.4.2",
        language = Migration.MigrationLanguage.JAVA,
        domain = Migration.MigrationDomain.CATALOG,
        date = 20220725)
public class RecoverProbandSamplesInCases extends MigrationTool {

    @Override
    protected void run() throws Exception {
        MongoCollection<Document> audit = getMongoCollection(MongoDBAdaptorFactory.AUDIT_COLLECTION);
        ClinicalAnalysisConverter converter = new ClinicalAnalysisConverter();

        queryMongo(MongoDBAdaptorFactory.CLINICAL_ANALYSIS_COLLECTION,
                new Document()
                        .append("proband.samples", null)
                        .append("locked", false),
                Projections.include("uid", "id", "uuid", "studyUid", "proband"), (doc) -> {
                    ClinicalAnalysis currentCase = converter.convertToDataModelType(doc);
                    logger.info("Trying to recover Clinical Analysis [id: {}, uid: {}, uuid: {}]", currentCase.getId(),
                            currentCase.getUid(), currentCase.getUuid());

                    if (currentCase.getProband() == null) {
                        logger.warn("Proband not defined in ClinicalAnalysis '{}'. Skipping...", currentCase.getId());
                        return;
                    }

                    ClinicalAnalysis clinicalAnalysis = findUpdatedCaseInAudit(currentCase.getUuid(), currentCase.getProband().getId(),
                            audit, converter);
                    if (clinicalAnalysis == null) {
                        clinicalAnalysis = findCreatedCaseInAudit(currentCase.getUuid(), currentCase.getProband().getId(), audit,
                                converter);
                    }
                    if (clinicalAnalysis == null) {
                        throw new RuntimeException("ClinicalAnalysis '" + currentCase.getId() + "' not found in the audit collection.");
                    }
                    if (clinicalAnalysis.getProband() == null) {
                        throw new RuntimeException("Proband from ClinicalAnalysis '" + currentCase.getId()
                                +  "' could not be found in the audit collection.");
                    }
                    if (CollectionUtils.isEmpty(clinicalAnalysis.getProband().getSamples())) {
                        throw new RuntimeException("List of samples from proband of ClinicalAnalysis '" + currentCase.getId()
                                + "' could not be found in the audit collection.");
                    }

                    // The proband from the audit matches the one used in the case currently :)
                    Individual proband;
                    try {
                        proband = dbAdaptorFactory.getCatalogIndividualDBAdaptor().get(currentCase.getProband().getUid(),
                                new QueryOptions(QueryOptions.INCLUDE, IndividualDBAdaptor.QueryParams.SAMPLES.key())).first();
                    } catch (CatalogException e) {
                        throw new RuntimeException(e);
                    }
                    List<Sample> sampleList = new ArrayList<>(clinicalAnalysis.getProband().getSamples().size());

                    Map<String, Sample> sampleMap = new HashMap<>();
                    for (Sample sample : proband.getSamples()) {
                        sampleMap.put(sample.getId(), sample);
                    }
                    for (Sample sample : clinicalAnalysis.getProband().getSamples()) {
                        if (sampleMap.containsKey(sample.getId())) {
                            sampleList.add(sample);
                        } else {
                            logger.warn("Sample '{}' is no longer associated to the individual '{}'", sample.getId(), proband.getId());
                        }
                    }

                    if (sampleList.isEmpty()) {
                        logger.warn("The list of samples associated to the proband '{}' is empty after processing", proband.getId());
                    }
                    proband.setSamples(sampleList);

                    ObjectMap params = new ObjectMap(ClinicalAnalysisDBAdaptor.QueryParams.PROBAND.key(), proband);
                    try {
                        dbAdaptorFactory.getClinicalAnalysisDBAdaptor().update(currentCase.getUid(), params, null, QueryOptions.empty());
                    } catch (CatalogException e) {
                        throw new RuntimeException(e);
                    }
                });
    }

    private ClinicalAnalysis findCreatedCaseInAudit(String uuid, String probandId, MongoCollection<Document> audit,
                                                    ClinicalAnalysisConverter converter) {
        // Fetch case from Audit collection
        try (MongoCursor<Document> iterator = audit.find(
                Filters.and(
                        Filters.eq("resource", Enums.Resource.CLINICAL_ANALYSIS),
                        Filters.eq("resourceUuid", uuid),
                        Filters.eq("action", Enums.Action.CREATE)
                )
        ).iterator()) {
            if (iterator.hasNext()) {
                Document auditDoc = iterator.next();
                Document auditParams = auditDoc.get("params", Document.class);
                Document caseDoc = auditParams.get("clinicalAnalysis", Document.class);
                ClinicalAnalysis clinicalAnalysis = converter.convertToDataModelType(caseDoc);
                if (!probandId.equals(clinicalAnalysis.getProband().getId())) {
                    logger.error("Proband '{}' from case '{}' does not match the proband '{}' used on create.", probandId, uuid,
                            clinicalAnalysis.getProband().getId());
                } else {
                    logger.debug("Found case in Audit create");
                    return clinicalAnalysis;
                }
            }
        }

        return null;
    }

    private ClinicalAnalysis findUpdatedCaseInAudit(String uuid, String probandId, MongoCollection<Document> audit,
                                                    ClinicalAnalysisConverter converter) {
        // Fetch case from Audit collection
        try (MongoCursor<Document> iterator = audit.find(
                Filters.and(
                        Filters.eq("resource", Enums.Resource.CLINICAL_ANALYSIS),
                        Filters.eq("resourceUuid", uuid),
                        Filters.eq("action", Enums.Action.UPDATE)
                )
        ).sort(Sorts.descending("date")).iterator()) {
            while (iterator.hasNext()) {
                Document auditDoc = iterator.next();
                Document auditParams = auditDoc.get("params", Document.class);
                Document caseDoc = auditParams.get("clinicalAnalysis", Document.class);
                ClinicalAnalysis clinicalAnalysis = converter.convertToDataModelType(caseDoc);
                if (clinicalAnalysis.getProband() != null && probandId.equals(clinicalAnalysis.getProband().getId())
                        && CollectionUtils.isNotEmpty(clinicalAnalysis.getProband().getSamples())) {
                    logger.debug("Found case in Audit update");
                    return clinicalAnalysis;
                }
            }
        }

        return null;
    }


}

@@ -6306,32 +6306,39 @@
 LINICAL_ANALYSIS
+.name()
 ),%0A             
@@ -6446,16 +6446,23 @@
 n.CREATE
+.name()
 )%0A      
@@ -7749,16 +7749,23 @@
 ANALYSIS
+.name()
 ),%0A     
@@ -7881,16 +7881,23 @@
 n.UPDATE
+.name()
 )%0A      

fix migration using enum instead of string, #TASK-1470