package ru.yandex.canvas.repository.mongo;

import java.util.List;
import java.util.Optional;
import java.util.function.Supplier;

import javax.annotation.ParametersAreNonnullByDefault;

import com.mongodb.client.MongoCollection;
import com.mongodb.client.model.UpdateOptions;
import com.mongodb.client.result.UpdateResult;
import one.util.streamex.StreamEx;
import org.bson.Document;
import org.springframework.data.mongodb.core.MongoOperations;
import org.springframework.data.mongodb.core.convert.QueryMapper;
import org.springframework.data.mongodb.core.convert.UpdateMapper;
import org.springframework.data.mongodb.core.query.Query;
import org.springframework.data.mongodb.core.query.Update;
import org.springframework.data.util.CloseableIterator;

import ru.yandex.canvas.model.video.Addition;
import ru.yandex.canvas.repository.RepositoryUtils;
import ru.yandex.direct.tracing.Trace;
import ru.yandex.direct.tracing.TraceProfile;

@ParametersAreNonnullByDefault
public class MongoOperationsWrapper {

    private final MongoOperations mongoOperations;
    private final String repositoryTag;
    private final QueryMapper queryMapper;
    private final UpdateMapper updateMapper;

    public MongoOperationsWrapper(MongoOperations mongoOperations, String repositoryTag) {
        this.mongoOperations = mongoOperations;
        this.repositoryTag = repositoryTag;
        this.queryMapper = new QueryMapper(mongoOperations.getConverter());
        this.updateMapper = new UpdateMapper(mongoOperations.getConverter());
    }

    public <EntityType> List<EntityType> find(Query query, Class<EntityType> entityClass) {
        return databaseWrapper(() -> mongoOperations.find(query, entityClass), "find");
    }

    public <EntityType> EntityType findOne(Query query, Class<EntityType> entityClass) {
        return databaseWrapper(() -> mongoOperations.findOne(query, entityClass), "find");
    }

    public <EntityType> UpdateResult updateFirst(Query query, Update update, Class<EntityType> entityClass) {
        return databaseWrapper(() -> mongoOperations.updateFirst(query, update, entityClass), "update");
    }

    public <EntityType> UpdateResult updateMulti(Query query, Update update, Class<EntityType> entityClass) {
        return databaseWrapper(() -> mongoOperations.updateMulti(query, update, entityClass), "update");
    }

    public <EntityType> void insert(EntityType entity) {
        try (TraceProfile profile = Trace.current().profile("db:insert", repositoryTag)) {
            mongoOperations.insert(entity);
        }
    }

    public <EntityType extends Addition> List<EntityType> insertWithDups(List<EntityType> entities,
                                                                         Class<EntityType> entityClass) {
        return databaseWrapper(() -> RepositoryUtils.insertWithDups(entities, mongoOperations, entityClass), "insert");
    }

    public <EntityType> long count(Query query, Class<EntityType> entityClass) {
        return databaseWrapper(() -> mongoOperations.count(query, entityClass), "count");
    }

    public <EntityType> CloseableIterator<EntityType> stream(Query query, Class<EntityType> entityClass) {
        return databaseWrapper(() -> mongoOperations.stream(query, entityClass), "stream");
    }

    /**
     * Метод обновляет вложенные списки внутри коллекций.
     * Работать может неспешно, использовать без необходимости либо без хорошего первичного фильтра
     * (параметр query) не рекомендуетя
     *
     * @param query        Первичный фильтр для апдейта
     * @param update       Апдейт с плейсхолдером для arrayFilters
     * @param arrayFilters Фильтры для вложенного массива
     */
    public <EntityType> UpdateResult updateNestedArrays(Query query,
                                                        Update update,
                                                        List<Query> arrayFilters,
                                                        Class<EntityType> entityClass) {
        return databaseWrapper(() -> {
            String collectionName = entityClass
                    .getAnnotation(org.springframework.data.mongodb.core.mapping.Document.class)
                    .collection();
            MongoCollection<Document> collection = mongoOperations.getCollection(collectionName);
            Document queryDoc = queryMapper.getMappedObject(query.getQueryObject(), Optional.empty());
            Document updateDoc = updateMapper.getMappedObject(update.getUpdateObject(), Optional.empty());
            List<Document> arrayFiltersDoc = StreamEx.of(arrayFilters)
                    .map(filter -> queryMapper.getMappedObject(filter.getQueryObject(), Optional.empty()))
                    .toList();
            UpdateOptions updateOptions = new UpdateOptions().arrayFilters(arrayFiltersDoc);
            return collection.updateMany(queryDoc, updateDoc, updateOptions);
        }, "updateNestedArray");
    }

    private <REC> REC databaseWrapper(Supplier<REC> wrapper, String opName) {
        try (TraceProfile profile = Trace.current().profile("db:" + opName, repositoryTag)) {
            return wrapper.get();
        }
    }
}
