<commit_before>package io.stardog.stardao.mongodb;

import com.google.common.collect.ImmutableList;
import com.google.common.collect.ImmutableSet;
import com.mongodb.client.FindIterable;
import com.mongodb.client.MongoCollection;
import com.mongodb.client.model.IndexModel;
import io.stardog.stardao.core.AbstractDao;
import io.stardog.stardao.core.Results;
import io.stardog.stardao.core.Update;
import io.stardog.stardao.core.field.Field;
import io.stardog.stardao.core.field.FieldData;
import io.stardog.stardao.exceptions.DataNotFoundException;
import io.stardog.stardao.mongodb.mapper.DocumentMapper;
import io.stardog.stardao.mongodb.mapper.JacksonDocumentMapper;
import org.bson.Document;
import org.bson.types.ObjectId;

import java.sql.Date;
import java.time.Instant;
import java.time.LocalDate;
import java.util.*;

public abstract class AbstractMongoDao<M,P,K,I> extends AbstractDao<M,P,K,I> {
    private final MongoCollection<Document> collection;
    private final DocumentMapper<M> modelMapper;
    private final DocumentMapper<P> partialMapper;
    public static final String ID_FIELD = "_id";

    public AbstractMongoDao(Class<M> modelClass, Class<P> partialClass, MongoCollection<Document> collection) {
        super(modelClass, partialClass);
        this.collection = collection;
        this.modelMapper = new JacksonDocumentMapper<>(modelClass,
                JacksonDocumentMapper.DEFAULT_OBJECT_MAPPER,
                getFieldData());
        this.partialMapper = new JacksonDocumentMapper<>(partialClass,
                JacksonDocumentMapper.DEFAULT_OBJECT_MAPPER,
                getFieldData());
    }

    public AbstractMongoDao(Class<M> modelClass, Class<P> partialClass, MongoCollection<Document> collection, DocumentMapper<M> modelMapper, DocumentMapper<P> partialMapper) {
        super(modelClass, partialClass);
        this.collection = collection;
        this.modelMapper = modelMapper;
        this.partialMapper = partialMapper;
    }

    /**
     * MongoDB always stores the id field as _id, so any id field should automatically get assigned @StoragePath("_id")
     * @return FieldData
     */
    @Override
    protected FieldData generateFieldData() {
        FieldData fieldData = super.generateFieldData();
        Field id = fieldData.getId();
        if (id != null) {
            id = id.toBuilder().storageName(ID_FIELD).build();
            Map<String,Field> fieldMap = new HashMap<>();
            fieldMap.putAll(fieldData.getMap());
            fieldMap.put(id.getName(), id);
            fieldData = fieldData.toBuilder().id(id).map(fieldMap).build();
        }
        return fieldData;
    }

    public MongoCollection<Document> getCollection() {
        return collection;
    }

    public String getCollectionName() {
        return collection.getNamespace().getCollectionName();
    }

    public DocumentMapper<M> getModelMapper() {
        return modelMapper;
    }

    public DocumentMapper<P> getPartialMapper() {
        return partialMapper;
    }

    protected Object generateId() {
        return new ObjectId();
    }

    @Override
    public Optional<M> loadOpt(K id) {
        Document query = new Document(ID_FIELD, id);
        Document doc = collection.find(query).limit(1).first();
        return Optional.ofNullable(modelMapper.toObject(doc));
    }

    protected M loadByQuery(Document query) {
        Document doc = collection.find(query).limit(1).first();
        if (doc == null) {
            throw new DataNotFoundException(getDisplayModelName() + " not found");
        }
        return modelMapper.toObject(doc);
    }

    protected Optional<M> loadByQueryOpt(Document query) {
        Document doc = collection.find(query).limit(1).first();
        if (doc == null) {
            return Optional.empty();
        }
        return Optional.of(modelMapper.toObject(doc));
    }

    /**
     * Given a query and a sort, return all documents that match the query in sorted order.
     * @param query query
     * @param sort  sort order
     * @return  results
     */
    protected Results<M,K> findByQuery(Document query, Document sort) {
        FindIterable<Document> iterable = getCollection().find(query);
        if (sort != null) {
            iterable = iterable.sort(sort);
        }
        List<M> result = new ArrayList<>();
        for (Document doc : iterable) {
            result.add(getModelMapper().toObject(doc));
        }
        return Results.of(result);
    }

    /**
     * Given a FindIterable query, add skip and limit to the query and return a resultset. The results will have
     * a next integer as the next "skip" value to use, or empty optional if we've exhausted results.
     * @param iterable  query<commit_msg>adding sort parameters to loadByQuery and loadByQueryOpt<commit_after>package io.stardog.stardao.mongodb;

import com.google.common.collect.ImmutableList;
import com.google.common.collect.ImmutableSet;
import com.mongodb.client.FindIterable;
import com.mongodb.client.MongoCollection;
import com.mongodb.client.model.IndexModel;
import io.stardog.stardao.core.AbstractDao;
import io.stardog.stardao.core.Results;
import io.stardog.stardao.core.Update;
import io.stardog.stardao.core.field.Field;
import io.stardog.stardao.core.field.FieldData;
import io.stardog.stardao.exceptions.DataNotFoundException;
import io.stardog.stardao.mongodb.mapper.DocumentMapper;
import io.stardog.stardao.mongodb.mapper.JacksonDocumentMapper;
import org.bson.Document;
import org.bson.types.ObjectId;

import java.sql.Date;
import java.time.Instant;
import java.time.LocalDate;
import java.util.*;

public abstract class AbstractMongoDao<M,P,K,I> extends AbstractDao<M,P,K,I> {
    private final MongoCollection<Document> collection;
    private final DocumentMapper<M> modelMapper;
    private final DocumentMapper<P> partialMapper;
    public static final String ID_FIELD = "_id";

    public AbstractMongoDao(Class<M> modelClass, Class<P> partialClass, MongoCollection<Document> collection) {
        super(modelClass, partialClass);
        this.collection = collection;
        this.modelMapper = new JacksonDocumentMapper<>(modelClass,
                JacksonDocumentMapper.DEFAULT_OBJECT_MAPPER,
                getFieldData());
        this.partialMapper = new JacksonDocumentMapper<>(partialClass,
                JacksonDocumentMapper.DEFAULT_OBJECT_MAPPER,
                getFieldData());
    }

    public AbstractMongoDao(Class<M> modelClass, Class<P> partialClass, MongoCollection<Document> collection, DocumentMapper<M> modelMapper, DocumentMapper<P> partialMapper) {
        super(modelClass, partialClass);
        this.collection = collection;
        this.modelMapper = modelMapper;
        this.partialMapper = partialMapper;
    }

    /**
     * MongoDB always stores the id field as _id, so any id field should automatically get assigned @StoragePath("_id")
     * @return FieldData
     */
    @Override
    protected FieldData generateFieldData() {
        FieldData fieldData = super.generateFieldData();
        Field id = fieldData.getId();
        if (id != null) {
            id = id.toBuilder().storageName(ID_FIELD).build();
            Map<String,Field> fieldMap = new HashMap<>();
            fieldMap.putAll(fieldData.getMap());
            fieldMap.put(id.getName(), id);
            fieldData = fieldData.toBuilder().id(id).map(fieldMap).build();
        }
        return fieldData;
    }

    public MongoCollection<Document> getCollection() {
        return collection;
    }

    public String getCollectionName() {
        return collection.getNamespace().getCollectionName();
    }

    public DocumentMapper<M> getModelMapper() {
        return modelMapper;
    }

    public DocumentMapper<P> getPartialMapper() {
        return partialMapper;
    }

    protected Object generateId() {
        return new ObjectId();
    }

    @Override
    public Optional<M> loadOpt(K id) {
        Document query = new Document(ID_FIELD, id);
        Document doc = collection.find(query).limit(1).first();
        return Optional.ofNullable(modelMapper.toObject(doc));
    }

    protected M loadByQuery(Document query) {
        return loadByQuery(query, null);
    }

    protected M loadByQuery(Document query, Document sort) {
        FindIterable<Document> find = collection.find(query);
        if (sort != null) {
            find.sort(sort);
        }
        Document doc = find.limit(1).first();
        if (doc == null) {
            throw new DataNotFoundException(getDisplayModelName() + " not found");
        }
        return modelMapper.toObject(doc);
    }

    protected Optional<M> loadByQueryOpt(Document query) {
        return loadByQueryOpt(query, null);
    }

    protected Optional<M> loadByQueryOpt(Document query, Document sort) {
        FindIterable<Document> find = collection.find(query);
        if (sort != null) {
            find.sort(sort);
        }
        Document doc = find.limit(1).first();
        if (doc == null) {
            return Optional.empty();
        }
        return Optional.of(modelMapper.toObject(doc));
    }

    /**
     * Given a query and a sort, return all documents that match the query in sorted order.
     * @param query query
     * @param sort  sort order
     * @return  results
     */
    protected Results<M,K> findByQuery(Document query, Document sort) {
        FindIterable<Document> iterable = getCollection().find(query);
        if (sort != null) {
            iterable = iterable.sort(sort);
        }
        List<M> result = new ArrayList<>();
        for (Document doc : iterable) {
            result.add(getModelMapper().toObject(doc));
        }
        return Results.of(result);
    }

    /**
     * Given a FindIterable query, add skip and limit to the query and return a resultset. The results will have
     * a next integer as the next "skip" value to use, or empty optional if we've exhausted results.
     * @param iterable  query
{'fraction_non_alphanumeric': 0.10391430573630994, 'fraction_numerical': 0.0006006607267994795, 'mean_word_length': 2.994402239104358, 'pattern_counts': {'":': 0, '<': 58, '<?xml version=': 0, '>': 58, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '142609', 'n_tokens_mistral': 2809, 'n_tokens_neox': 2762, 'n_words': 657}