package com.wxz.mongodb;

import com.alibaba.fastjson.JSON;
import com.mongodb.BasicDBObject;
import com.mongodb.MongoClient;
import com.mongodb.ServerAddress;
import com.mongodb.client.FindIterable;
import com.mongodb.client.MongoCollection;
import com.mongodb.client.MongoCursor;
import com.mongodb.client.MongoDatabase;
import com.mongodb.client.model.IndexOptions;
import com.mongodb.client.model.UpdateOptions;
import com.mongodb.client.result.DeleteResult;
import com.mongodb.client.result.UpdateResult;
import org.bson.Document;

import java.text.MessageFormat;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.function.Consumer;
import java.util.stream.Collectors;

/**
 * @author [email protected]
 * @type class
 * @date 2017/9/21 -16:54
 */
public class MongodbDataAccess {
    private MongoClient sMongoClient;
    private MongoDatabase sMongoDatabase;

    public MongodbDataAccess(String ip, int port, String databaseName) {
        sMongoClient = MongodbClientFactory.getClient(ip, port);
        sMongoDatabase = sMongoClient.getDatabase(databaseName);
    }


    /**
     * 集合满足查询条件结果总数
     *
     * @param collectionName
     * @param query
     * @return
     */
    public long count(String collectionName, MongodbQuery query) {
        MongoCollection<Document> collection = sMongoDatabase.getCollection(collectionName);
        if (query == null) {
            return collection.count();
        }
        return collection.count(query.getQuery());
    }

    /**
     * 集合总条数
     *
     * @param collectionName
     * @return
     */
    public long count(String collectionName) {
        return count(collectionName, null);
    }

    /**
     * 查询
     *
     * @param collectionName 集合名
     * @param query          查询条件
     * @param fields         返回字段或者排除字段
     * @param sort           排序方式
     * @param pageInfo       分页
     * @return
     */
    public List<Map<String, Object>> find(
            String collectionName,
            MongodbQuery query, MongodbFields fields,
            MongodbSort sort, MongodbPageInfo pageInfo) {
        List<Map<String, Object>> resultMapList = new ArrayList<Map<String, Object>>();
        MongoCollection<Document> collection = sMongoDatabase.getCollection(collectionName);
        FindIterable<Document> findIterable = collection.find(query == null ? new Document() : query.getQuery());
        if (fields == null) {
            findIterable.projection(new MongodbFields().getDbObject());
        } else {
            findIterable.projection(fields.getDbObject());
        }
        if (sort != null) {
            findIterable.sort(sort.getDbObject());
        }
        if (pageInfo != null) {
            int startPos = pageInfo.getPageIndex() * pageInfo.getPageSize();
            int rows = pageInfo.getPageSize();
            if (startPos > 0) {
                findIterable.skip(startPos - 1);
            }
            findIterable.limit(rows);
        }
        MongoCursor<Document> cursor = findIterable.iterator();
        try {
            while (cursor.hasNext()) {
                Map<String, Object> document = cursor.next();
                resultMapList.add(document);
            }
        } finally {
            cursor.close();
        }

        return resultMapList;
    }

    /**
     * 查询
     *
     * @param collectionName 集合名
     * @param query          查询条件
     * @param fields         返回字段或者排除字段
     * @param sort           排序
     * @return
     */
    public List<Map<String, Object>> find(
            String collectionName,
            MongodbQuery query, MongodbFields fields,
            MongodbSort sort) {
        return find(collectionName, query, fields, sort, null);
    }

    /**
     * 查询并逐条处理
     *
     * @param collectionName 集合名
     * @param query          查询条件
     * @param fields         返回字段或者排除字段
     * @param sort           排序方式
     * @param consumer       记录处理
     * @return
     */
    public void findAndConsumer(
            String collectionName,
            MongodbQuery query, MongodbFields fields,
            MongodbSort sort, Consumer<Map<String, Object>> consumer) {
        MongoCollection<Document> collection = sMongoDatabase.getCollection(collectionName);
        FindIterable<Document> findIterable = collection.find(query == null ? new Document() : query.getQuery());
        if (fields == null) {
            findIterable.projection(new MongodbFields().getDbObject());
        } else {
            findIterable.projection(fields.getDbObject());
        }
        if (sort != null) {
            findIterable.sort(sort.getDbObject());
        }
        MongoCursor<Document> cursor = findIterable.iterator();
        try {
            while (cursor.hasNext()) {
                Map<String, Object> document = cursor.next();
                consumer.accept(document);
            }
        } finally {
            cursor.close();
        }
    }


    /**
     * 查询
     *
     * @param collectionName 集合名
     * @param query          查询条件
     * @param fields         返回字段或者排除字段
     * @return
     */
    public List<Map<String, Object>> find(
            String collectionName,
            MongodbQuery query, MongodbFields fields) {
        return find(collectionName, query, fields, null, null);
    }


    /**
     * 查询
     *
     * @param clazz          类
     * @param collectionName 集合名
     * @param query          查询条件
     * @param fields         返回字段或者排除字段
     * @param sort           排序方式
     * @param <T>
     * @return
     */
    public <T> List<T> find(
            Class<T> clazz, String collectionName,
            MongodbQuery query, MongodbFields fields,
            MongodbSort sort, MongodbPageInfo pageInfo) {
        List<T> resultMapList = new ArrayList<T>();
        MongoCollection<Document> collection = sMongoDatabase.getCollection(collectionName);
        FindIterable<Document> findIterable = collection.find(query == null ? new Document() : query.getQuery());
        if (fields == null) {
            findIterable.projection(new MongodbFields().getDbObject());
        } else {
            findIterable.projection(fields.getDbObject());
        }
        if (sort != null) {
            findIterable.sort(sort.getDbObject());
        }
        if (pageInfo != null) {
            int startPos = pageInfo.getPageIndex() * pageInfo.getPageSize();
            int rows = pageInfo.getPageSize();
            if (startPos > 0) {
                findIterable.skip(startPos - 1);
            }
            findIterable.limit(rows);
        }
        MongoCursor<Document> cursor = findIterable.iterator();
        try {
            while (cursor.hasNext()) {
                Document document = cursor.next();
                T model = JSON.parseObject(JSON.toJSONString(document), clazz);
                resultMapList.add(model);
            }
        } finally {
            cursor.close();
        }

        return resultMapList;
    }

    /**
     * 查询并逐条处理
     *
     * @param clazz          类
     * @param collectionName 集合名
     * @param query          查询条件
     * @param fields         返回字段或者排除字段
     * @param sort           排序方式
     * @param consumer       记录处理
     * @param <T>
     * @return
     */
    public <T> void findAndConsumer(
            Class<T> clazz, String collectionName,
            MongodbQuery query, MongodbFields fields,
            MongodbSort sort, Consumer<T> consumer) {
        MongoCollection<Document> collection = sMongoDatabase.getCollection(collectionName);
        FindIterable<Document> findIterable = collection.find(query == null ? new Document() : query.getQuery());
        if (fields == null) {
            findIterable.projection(new MongodbFields().getDbObject());
        } else {
            findIterable.projection(fields.getDbObject());
        }
        if (sort != null) {
            findIterable.sort(sort.getDbObject());
        }
        MongoCursor<Document> cursor = findIterable.iterator();
        try {
            while (cursor.hasNext()) {
                Document document = cursor.next();
                T model = JSON.parseObject(JSON.toJSONString(document), clazz);
                consumer.accept(model);
            }
        } finally {
            cursor.close();
        }

    }

    /**
     * @param clazz          类
     * @param collectionName 集合名
     * @param query          查询条件
     * @param fields         返回字段或者排除字段
     * @param sort           排序
     * @param <T>
     * @return
     */
    public <T> List<T> find(
            Class<T> clazz, String collectionName,
            MongodbQuery query, MongodbFields fields,
            MongodbSort sort) {
        return find(clazz, collectionName, query, fields, sort, null);
    }


    /**
     * 查询
     *
     * @param clazz          类
     * @param collectionName 集合名
     * @param query          查询条件
     * @param fields         返回字段或者排除字段
     * @param <T>
     * @return
     */
    public <T> List<T> find(
            Class<T> clazz, String collectionName,
            MongodbQuery query, MongodbFields fields) {
        return find(clazz, collectionName, query, fields, null, null);
    }


    /**
     * 查询
     *
     * @param clazz          类
     * @param collectionName 集合名
     * @param <T>
     * @return
     */
    public <T> List<T> findAll(Class<T> clazz, String collectionName) {
        return findAll(clazz, collectionName, null);
    }

    /**
     * 查询
     *
     * @param clazz          类
     * @param collectionName 集合名
     * @param sort           排序
     * @param <T>
     * @return
     */
    public <T> List<T> findAll(Class<T> clazz, String collectionName, MongodbSort sort) {
        List<T> resultMapList = new ArrayList<T>();
        MongoCollection<Document> collection = sMongoDatabase.getCollection(collectionName);
        FindIterable<Document> findIterable = collection.find();
        if(sort != null) {
            findIterable.sort(sort.getDbObject());
        }
        MongoCursor<Document> cursor = findIterable.iterator();
        try {
            while (cursor.hasNext()) {
                Document document = cursor.next();
                T parseObject = JSON.parseObject(JSON.toJSONString(document), clazz);
                resultMapList.add(parseObject);
            }
        } finally {
            cursor.close();
        }

        return resultMapList;
    }

    /**
     * 查询
     *
     * @param collectionName 集合名
     * @param sort           排序方式
     * @return
     */
    public List<Map<String, Object>> findAll(String collectionName, MongodbSort sort) {
        List<Map<String, Object>> resultMapList = new ArrayList<>();
        MongoCollection<Document> collection = sMongoDatabase.getCollection(collectionName);
        FindIterable<Document> findIterable = collection.find();
        if (sort != null) {
            findIterable.sort(sort.getDbObject());
        }
        MongoCursor<Document> cursor = findIterable.iterator();
        try {
            while (cursor.hasNext()) {
                Document document = cursor.next();
                resultMapList.add(document);
            }
        } finally {
            cursor.close();
        }

        return resultMapList;
    }

    /**
     * 查询
     *
     * @param collectionName 集合名
     * @return
     */
    public List<Map<String, Object>> findAll(String collectionName) {
        return findAll(collectionName, null);
    }

    /**
     * 查询一个
     *
     * @param collectionName 集合名
     * @param query          查询条件
     * @param fields         返回字段或者排除字段
     * @param sort
     * @return
     */
    public Map<String, Object> findOne(
            String collectionName,
            MongodbQuery query, MongodbFields fields, MongodbSort sort) {
        MongoCollection<Document> collection = sMongoDatabase.getCollection(collectionName);
        FindIterable<Document> findIterable = collection.find(query.getQuery());
        if (fields == null) {
            findIterable.projection(new MongodbFields().getDbObject());
        } else {
            findIterable.projection(fields.getDbObject());
        }
        if (sort != null) {
            findIterable.sort(sort.getDbObject());
        }
        findIterable.limit(1);
        MongoCursor<Document> cursor = findIterable.iterator();
        try {
            if (cursor.hasNext()) {
                return cursor.next();
            }
        } finally {
            cursor.close();
        }
        return null;
    }

    /**
     * 查询一个
     *
     * @param clazz          类
     * @param collectionName 集合名
     * @param query          查询条件
     * @param fields         返回字段或者排除字段
     * @param sort
     * @param <T>
     * @return
     */
    public <T> T findOne(
            Class<T> clazz, String collectionName,
            MongodbQuery query, MongodbFields fields, MongodbSort sort) {
        MongoCollection<Document> collection = sMongoDatabase.getCollection(collectionName);
        FindIterable<Document> findIterable = collection.find(query.getQuery());
        if (fields == null) {
            findIterable.projection(new MongodbFields().getDbObject());
        } else {
            findIterable.projection(fields.getDbObject());
        }
        if (sort != null) {
            findIterable.sort(sort.getDbObject());
        }
        findIterable.limit(1);
        MongoCursor<Document> cursor = findIterable.iterator();
        try {
            if (cursor.hasNext()) {
                Document document = cursor.next();
                return JSON.parseObject(document.toJson(), clazz);
            }
        } finally {
            cursor.close();
        }
        return null;
    }

    /**
     * 新增或者更新
     *
     * @param collectionName 集合名
     * @param query          查询条件
     * @param descData       目标数据
     * @return
     */
    public boolean upsert(String collectionName, MongodbQuery query, Map<String, Object> descData) {
        MongoCollection collection = sMongoDatabase.getCollection(collectionName);
        UpdateOptions options = new UpdateOptions();
        options.upsert(true);
        BasicDBObject updateSetValue = new BasicDBObject("$set", descData);
        UpdateResult updateResult = collection.updateMany(query.getQuery(), updateSetValue, options);
        return updateResult.getUpsertedId() != null ||
                (updateResult.getMatchedCount() > 0 && updateResult.getModifiedCount() > 0);
    }

    /**
     * 更新符合条件的一条记录
     *
     * @param collectionName 集合名
     * @param query          查询条件
     * @param update         更新设置
     * @return
     */
    public boolean updateOne(String collectionName, MongodbQuery query, MongodbUpdates update) {
        MongoCollection collection = sMongoDatabase.getCollection(collectionName);
        UpdateResult updateResult = collection.updateOne(query.getQuery(), update.getUpdates());
        return updateResult.getUpsertedId() != null ||
                (updateResult.getMatchedCount() > 0 && updateResult.getModifiedCount() > 0);
    }

    /**
     * 更新符合条件的所有记录
     *
     * @param collectionName 集合名
     * @param query          查询条件
     * @param update         更新设置
     * @return
     */
    public boolean updateMany(String collectionName, MongodbQuery query, MongodbUpdates update) {
        MongoCollection collection = sMongoDatabase.getCollection(collectionName);
        UpdateResult updateResult = collection.updateMany(query.getQuery(), update.getUpdates());
        return updateResult.getUpsertedId() != null ||
                (updateResult.getMatchedCount() > 0 && updateResult.getModifiedCount() > 0);
    }

    /**
     * 移除指定字段
     *
     * @param collectionName
     * @param query
     * @param fieldName
     */
    public void removeField(String collectionName, MongodbQuery query, String fieldName) {
        MongoCollection collection = sMongoDatabase.getCollection(collectionName);
        UpdateOptions options = new UpdateOptions();
        options.upsert(true);
        collection.updateMany(query.getQuery(), MongodbUpdates.unset(fieldName).getUpdates(), options);
    }

    /**
     * 插入一条
     *
     * @param collectionName
     * @param data
     * @return
     */
    public boolean insertOne(String collectionName, Map<String, Object> data) {
        MongoCollection collection = sMongoDatabase.getCollection(collectionName);
        collection.insertOne(new Document(data));
        return true;
    }

    /**
     * 插入记录,不做唯一性检查
     *
     * @param collectionName
     * @param data
     * @param <T>
     * @return
     */
    public <T> boolean insertOne(String collectionName, T data) {
        MongoCollection collection = sMongoDatabase.getCollection(collectionName);
        collection.insertOne(Document.parse(JSON.toJSONString(data)));
        return true;
    }

    /**
     * 批量插入
     *
     * @param collectionName
     * @param dataList
     * @return
     */
    public boolean insertMany(String collectionName, List<Map<String, Object>> dataList) {
        MongoCollection collection = sMongoDatabase.getCollection(collectionName);
        List<Document> documentList = new ArrayList<>();
        for (Map<String, Object> data : dataList) {
            documentList.add(new Document(data));
        }
        collection.insertMany(documentList);
        return true;
    }

    /**
     * 批量插入
     *
     * @param collectionName
     * @param dataList
     * @return
     */
    public <T> boolean insertModelList(String collectionName, List<T> dataList) {
        MongoCollection collection = sMongoDatabase.getCollection(collectionName);
        List<Document> documentList = dataList.stream()
                .map(data -> Document.parse(JSON.toJSONString(data)))
                .collect(Collectors.toList());
        collection.insertMany(documentList);
        return true;
    }

    /**
     * 按条件删除
     *
     * @param collectionName 集合名
     * @param query          查询条件
     * @return
     */
    public boolean delete(String collectionName, MongodbQuery query) {
        MongoCollection collection = sMongoDatabase.getCollection(collectionName);
        DeleteResult deleteResult = collection.deleteMany(query.getQuery());
        return deleteResult.getDeletedCount() > 0;
    }

    /**
     * 删除集合
     *
     * @param collectionName 集合名
     */
    public void dropCollection(String collectionName) {
        MongoCollection collection = sMongoDatabase.getCollection(collectionName);
        collection.drop();
    }

    /**
     * 创建索引
     *
     * @param collectionName 集合
     * @param indexName      索引名
     * @param keys           键信息,如 securityCode:1
     * @param isUnique       是否唯一索引
     */
    public void createIndex(String collectionName, String indexName, Map<String, Integer> keys, boolean isUnique) {
        MongoCollection collection = sMongoDatabase.getCollection(collectionName);
        IndexOptions indexOptions = new IndexOptions();
        indexOptions.background(true);
        indexOptions.unique(isUnique);
        indexOptions.name(indexName);
        Map<String, Object> keyIndexs = new HashMap<>(16);
        for (Map.Entry<String, Integer> entry : keys.entrySet()) {
            keyIndexs.put(entry.getKey(), entry.getValue() > 0 ? 1 : entry.getValue() == 0 ? 0 : -1);
        }
        collection.createIndex(new Document(keyIndexs), indexOptions);
    }

    /**
     * 删除索引
     *
     * @param collectionName 集合名
     * @param indexName      索引名
     */
    public void dropIndex(String collectionName, String indexName) {
        MongoCollection collection = sMongoDatabase.getCollection(collectionName);
        collection.dropIndex(indexName);
    }

    /**
     * 删除所有索引
     *
     * @param collectionName 集合名
     */
    public void dropIndexes(String collectionName) {
        MongoCollection collection = sMongoDatabase.getCollection(collectionName);
        collection.dropIndexes();
    }

    /**
     * 获取获取数据所在Mongo数据库信息,包括集合名
     *
     * @param collectionName 集合名
     * @return
     */
    public String getMongoInfo(String collectionName) {
        ServerAddress address = sMongoClient.getAddress();
        return MessageFormat.format(
                "Mongo[host={0},port={1},database={2},collectionName={3}]",
                address.getHost(),
                address.getPort(),
                sMongoDatabase.getName(),
                collectionName);
    }
}