package org.hisp.dhis.android.core.sms.data.localdbrepository.internal;

import org.hisp.dhis.android.core.common.State;
import org.hisp.dhis.android.core.dataset.DataSet;
import org.hisp.dhis.android.core.dataset.DataSetCollectionRepository;
import org.hisp.dhis.android.core.dataset.DataSetCompleteRegistration;
import org.hisp.dhis.android.core.dataset.DataSetCompleteRegistrationCollectionRepository;
import org.hisp.dhis.android.core.dataset.DataSetElement;
import org.hisp.dhis.android.core.dataset.internal.DataSetCompleteRegistrationStore;
import org.hisp.dhis.android.core.datavalue.DataValue;
import org.hisp.dhis.android.core.datavalue.DataValueCollectionRepository;
import org.hisp.dhis.android.core.datavalue.DataValueModule;
import org.hisp.dhis.android.core.datavalue.internal.DataValueStore;

import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;

import javax.inject.Inject;

import io.reactivex.Completable;
import io.reactivex.Single;

class DataSetsStore {
    private final DataValueModule dataValueModule;
    private final DataValueStore dataValueStore;
    private final DataSetCompleteRegistrationStore dataSetStore;
    private final DataSetCollectionRepository dataSetRepository;
    private final DataSetCompleteRegistrationCollectionRepository completeRegistrationRepository;

    @Inject
    DataSetsStore(DataValueModule dataValueModule,
                  DataValueStore dataValueStore,
                  DataSetCompleteRegistrationStore dataSetStore,
                  DataSetCollectionRepository dataSetRepository,
                  DataSetCompleteRegistrationCollectionRepository completeRegistrationRepository) {
        this.dataValueModule = dataValueModule;
        this.dataValueStore = dataValueStore;
        this.dataSetStore = dataSetStore;
        this.dataSetRepository = dataSetRepository;
        this.completeRegistrationRepository = completeRegistrationRepository;
    }

    Single<List<DataValue>> getDataValues(String dataSetUid, String orgUnit,
                                          String period, String attributeOptionComboUid) {
        return Single.fromCallable(() -> {
            DataSet dataSet = dataSetRepository
                    .byUid().eq(dataSetUid)
                    .withDataSetElements()
                    .one().blockingGet();

            List<String> dataElementUids = new ArrayList<>();
            if (dataSet != null && dataSet.dataSetElements() != null) {
                for (DataSetElement dataSetElement : dataSet.dataSetElements()) {
                    dataElementUids.add(dataSetElement.dataElement().uid());
                }
            }

            DataValueCollectionRepository baseDataValuesRepo = dataValueModule.dataValues()
                    .byDataElementUid().in(dataElementUids)
                    .byOrganisationUnitUid().eq(orgUnit)
                    .byPeriod().eq(period)
                    .byAttributeOptionComboUid().eq(attributeOptionComboUid);

            List<DataValue> dataValues = baseDataValuesRepo
                    .byState().in(Arrays.asList(State.uploadableStates()))
                    .blockingGet();

            // TODO Workaround to prevent empty lists. Not supported in compression library
            if (dataValues.isEmpty()) {
                List<DataValue> allDataValues = baseDataValuesRepo.blockingGet();

                if (!allDataValues.isEmpty()) {
                    dataValues = allDataValues.subList(0, 1);
                }
            }

            return dataValues;
        });
    }

    Completable updateDataSetValuesState(String dataSet,
                                         String orgUnit,
                                         String period,
                                         String attributeOptionComboUid,
                                         State state) {
        return getDataValues(dataSet, orgUnit, period, attributeOptionComboUid)
                .flattenAsObservable(items -> items)
                .flatMapCompletable(item -> Completable.fromAction(() ->
                        dataValueStore.setState(item, state))
                );
    }

    Completable updateDataSetCompleteRegistrationState(String dataSetId,
                                                       String orgUnit,
                                                       String period,
                                                       String attributeOptionComboUid,
                                                       State state) {
        return Completable.fromAction(() -> {
            DataSetCompleteRegistration dataSet = completeRegistrationRepository
                    .byDataSetUid().eq(dataSetId)
                    .byOrganisationUnitUid().eq(orgUnit)
                    .byPeriod().eq(period)
                    .byAttributeOptionComboUid().eq(attributeOptionComboUid)
                    .one().blockingGet();
            if (dataSet != null) {
                dataSetStore.setState(dataSet, state);
            }
        });
    }
}

@@ -118,130 +118,8 @@
 te;%0A
-import org.hisp.dhis.android.core.dataset.DataSet;%0Aimport org.hisp.dhis.android.core.dataset.DataSetCollectionRepository;%0A
 impo
@@ -185,16 +185,16 @@
 ration;%0A
+
 import o
@@ -280,66 +280,8 @@
 ry;%0A
-import org.hisp.dhis.android.core.dataset.DataSetElement;%0A
 impo
@@ -626,36 +626,8 @@
 e;%0A%0A
-import java.util.ArrayList;%0A
 impo
@@ -953,73 +953,8 @@
 re;%0A
-    private final DataSetCollectionRepository dataSetRepository;%0A
     
@@ -1229,73 +1229,8 @@
 re,%0A
-                  DataSetCollectionRepository dataSetRepository,%0A
     
@@ -1465,60 +1465,8 @@
 re;%0A
-        this.dataSetRepository = dataSetRepository;%0A
     
@@ -1761,512 +1761,8 @@
 %3E %7B%0A
-            DataSet dataSet = dataSetRepository%0A                    .byUid().eq(dataSetUid)%0A                    .withDataSetElements()%0A                    .one().blockingGet();%0A%0A            List%3CString%3E dataElementUids = new ArrayList%3C%3E();%0A            if (dataSet != null && dataSet.dataSetElements() != null) %7B%0A                for (DataSetElement dataSetElement : dataSet.dataSetElements()) %7B%0A                    dataElementUids.add(dataSetElement.dataElement().uid());%0A                %7D%0A            %7D%0A%0A
     
@@ -1880,39 +1880,25 @@
 Data
-ElementUid().in(dataElemen
+SetUid(dataSe
 tUid
-s
 )%0A  

Use byDataSetUid in SMS local repository