package lc.esp.eo;

import com.mongodb.*;
import com.mongodb.client.*;
import com.mongodb.client.gridfs.GridFSBucket;
import com.mongodb.client.gridfs.GridFSBuckets;
import com.mongodb.client.gridfs.model.GridFSUploadOptions;
import com.mongodb.client.model.UpdateOptions;
import com.mongodb.client.result.DeleteResult;
import com.mongodb.client.result.UpdateResult;
import lc.esp.eo.sdk.EO;
import lc.esp.eo.sdk.EOAssociation;
import lc.esp.eo.sdk.EOId;
import lc.esp.eo.sdk.schema.AssociationSchema;
import lc.esp.eo.sdk.schema.Model;
import lc.mecha.log.MechaLogger;
import lc.mecha.log.MechaLoggerFactory;
import org.bson.Document;
import org.bson.types.ObjectId;

import java.io.*;
import java.nio.file.Path;
import java.time.Instant;
import java.util.HashSet;
import java.util.Iterator;
import java.util.Map;

/**
 * This class implements a transaction with the back-end. This includes both functionality for the objects themselves
 * and EOFS.
 *
 * @author Alex Leigh
 * @since mk20 (TOKYO PROTOCOL)
 */
public class EOTransaction implements AutoCloseable {
    private static final MechaLogger logger = MechaLoggerFactory.getLogger(EOTransaction.class);
    private final ClientSession session;
    private final MongoClient client;
    private final String databaseName;
    private final MongoDatabase db;
    public final static String KEY_ID = "_id";
    public final static String KEY_DATA = "data";
    public final static String KEY_TYPE = "type";
    public final static String KEY_DSTOBJ = "dstObj";
    public final static String KEY_SRCOBJ = "srcObj";
    public final static String KEY_CONTENT_TYPE = "content_type";
    public final static String COLLECTION_OBJ = "eo_obj";
    public final static String COLLECTION_ASSOC = "eo_assoc";
    public final static String COLLECTION_CHUNKS = "chunks";
    public final static String COLLECTION_FILES = "files";
    public final static String KEY_FS_FILENAME = "name";
    private final static String TYPE_FS_LINK = "link";
    private final static String TYPE_FS_FILE = "file";
    private final static String KEY_FS_GRIDFS_ID = "gridfs_id";
    public final static String KEY_CTIME = "ctime";
    public final static String TYPE_INODE = "inode";
    public final static String KEY_NAME = "name";
    public static String KEY_FS_LINK = "link";
    public static String VALUE_FS_PARENT = "..";
    public static String VALUE_FS_SELF = ".";
    public static String KEY_FS_SEPARATOR = "/";
    private HashSet<ObjectId> pendingFiles;
    private final Model model;

    /**
     * Create a new transaction.
     */
    public EOTransaction(MongoClient client, String databaseName, Model model) {
        this.model = model;
        this.client = client;
        this.session = this.client.startSession();
        this.session.startTransaction(TransactionOptions.builder().writeConcern(WriteConcern.ACKNOWLEDGED).build());
        this.databaseName = "eo_" + databaseName;
        this.db = client.getDatabase(this.databaseName);
    }

    /**
     * Initialize a new database for use with EO. This method runs outside of the transaction.
     * <p>
     * Note that there is a propagation delay from the time this call is made to when it is consistent across the
     * cluster. If the caller immediately tries to perform a data operation on the collections it may result in
     * errors.
     */
    public void initialize() {
        shardByKey(COLLECTION_OBJ, KEY_ID);
        shardByKey(COLLECTION_ASSOC, KEY_SRCOBJ);
        shardByKey(COLLECTION_FILES, KEY_ID);
        shardByKey(COLLECTION_CHUNKS, KEY_ID);
    }

    /**
     * Shard the given collection to the given key using a hashed strategy. This call is asynchronous and may return
     * before all of the shards have been updated.
     */
    private void shardByKey(String collectionName, String keyName) {
        String path = databaseName + "." + collectionName;
        MongoDatabase adminDb = client.getDatabase("admin");
        Document cmd = new Document("shardCollection", path).append("key", new Document(keyName, "hashed"));
        Document response = adminDb.runCommand(cmd);
        logger.debug("shardBykey: {}", response);
        // TODO: Wait for `db.collection.getShardDistribution()`?
    }

    /**
     * Create an object of the given type and return the newly created id.
     */
    public EOId objAdd(String type) {
        return objAdd(type, null);
    }

    public EOId objAdd(String type, Map<String, Object> values) {
        return objAdd(type, values, Instant.now());
    }

    /**
     * Create a new object of the given type, with the given user-defined data.
     */
    public EOId objAdd(String type, Map<String, Object> values, Instant ctime) {
        // FIXME: Returning ObjectId here is a leaky abstraction
        assert (type != null);
        MongoCollection<Document> col = db.getCollection(COLLECTION_OBJ);
        Document d = new Document();
        if (values != null) d.append(KEY_DATA, values);
        d.append(KEY_TYPE, type);
        d.append(KEY_CTIME, ctime.toEpochMilli());
        col.insertOne(session, d);
        ObjectId id = d.getObjectId(KEY_ID);
        return id == null ? null : new EOBsonId(id);
    }

    /**
     * Create an association between the source and destination objects of the given type.
     */
    public String assocAdd(EOId srcObjectId, EOId dstObjectId, String type) {
        return assocAdd(srcObjectId, dstObjectId, type, null);
    }

    public String assocAdd(EOId srcObjectId, EOId dstObjectId, String type,
                           Map<String, Serializable> values) {
        return assocAdd(srcObjectId, dstObjectId, type, values, Instant.now());
    }

    /**
     * Create a one-way association between two objects. The association may optionally contain a user-defined
     * data.
     */
    public String assocAdd(EOId srcObjectId, EOId dstObjectId, String type,
                           Map<String, Serializable> values, Instant ctime) {
        assert (srcObjectId != null);
        assert (dstObjectId != null);
        assert (type != null);
        assert (ctime != null);

        // FIXME: Do not duplicate the association if it already exists

        MongoCollection<Document> col = db.getCollection(COLLECTION_ASSOC);
        Document d = new Document();
        if (values != null) d.append(KEY_DATA, values);
        d.append(KEY_TYPE, type);
        d.append(KEY_SRCOBJ, ((EOBsonId) srcObjectId).getObjectId());
        d.append(KEY_DSTOBJ, ((EOBsonId) dstObjectId).getObjectId());
        d.append(KEY_CTIME, ctime.toEpochMilli());
        col.insertOne(session, d);
        ObjectId id = d.getObjectId(KEY_ID);

        logger.debug("Added association: {} {} -> {} {}", type, srcObjectId, dstObjectId, values);

        if (model != null) {
            AssociationSchema as = model.getAssocSchema(type);
            if (as.getInverseType() != null) {
                d = new Document();
                if (values != null) d.append(KEY_DATA, values);
                d.append(KEY_TYPE, as.getInverseType());
                d.append(KEY_SRCOBJ, ((EOBsonId) dstObjectId).getObjectId());
                d.append(KEY_DSTOBJ, ((EOBsonId) srcObjectId).getObjectId());
                d.append(KEY_CTIME, ctime.toEpochMilli());
                col.insertOne(session, d);
            }
        }

        return id == null ? null : id.toString();
    }

    /**
     * Return a count of the associations originating from the given source object of the given type.
     */
    public long assocCount(EOId srcObjectId, String type) {
        MongoCollection<Document> col = db.getCollection(COLLECTION_ASSOC);
        return col.countDocuments(session, new Document()
                .append(KEY_SRCOBJ, ((EOBsonId) srcObjectId).getObjectId())
                .append(KEY_TYPE, type));
    }

    /**
     * Delete the given association. If a model is defined in the transaction and an inverse association is defined,
     * then the inverse will also be deleted. The method returns the count of the associations which were deleted.
     */
    public long assocDelete(EOId srcObjectId, EOId dstObjectId, String type) {
        MongoCollection<Document> col = db.getCollection(COLLECTION_ASSOC);

        long cnt = 0;

        // deleteOne() under the premise that they are unique
        DeleteResult result = col.deleteOne(session, new Document()
                .append(KEY_TYPE, type)
                .append(KEY_SRCOBJ, ((EOBsonId) srcObjectId).getObjectId())
                .append(KEY_DSTOBJ, ((EOBsonId) dstObjectId).getObjectId()));

        cnt += result.getDeletedCount();

        if (model != null) {
            AssociationSchema as = model.getAssocSchema(type);
            if (as != null) {
                result = col.deleteOne(session, new Document()
                        .append(KEY_TYPE, as.getInverseType())
                        .append(KEY_SRCOBJ, ((EOBsonId) dstObjectId).getObjectId())
                        .append(KEY_DSTOBJ, ((EOBsonId) srcObjectId).getObjectId()));
                cnt += result.getDeletedCount();
            }
        }

        return cnt;
    }

    /**
     * Read a given object and return it to the caller in the form of an {@link EO}. Reading the object will not
     * query the underlying associations, so only the id, type, ctime, and any user-defined data is returned.
     */
    public EO objGet(EOId id) {
        MongoCollection<Document> col = db.getCollection(COLLECTION_OBJ);
        BasicDBObject query = new BasicDBObject();
        query.put(KEY_ID, ((EOBsonId) id).getObjectId());
        Document doc = col.find(session, query).first();

        if (doc == null) return null;
        return docToObject(doc);
    }

    /**
     * Update the object at the given id with the given data values. If <code>merged</code> is <code>true</code>
     * then the provided value will be merged with the existing object, otherwise, the data values in the object
     * will be cleared and replaced with the provided values.
     */
    public long objUpdate(EOId id, Map<String, Object> updates, boolean merge) {
        MongoCollection<Document> col = db.getCollection(COLLECTION_OBJ);

        BasicDBObject query = new BasicDBObject();
        query.put(KEY_ID, ((EOBsonId) id).getObjectId());

        BasicDBObject update = new BasicDBObject();
        if (!merge) {
            update.append("$set", new BasicDBObject(KEY_DATA, updates));
        } else {
            BasicDBObject set = new BasicDBObject();
            for (Map.Entry<String, Object> entry : updates.entrySet()) {
                set.put(KEY_DATA + "." + entry.getKey(), entry.getValue());
            }
            update.append("$set", set);
        }

        logger.debug("objUpdate: [query: {}] [update: {}]", query, update);

        UpdateOptions options = new UpdateOptions().upsert(false);
        UpdateResult result = col.updateOne(session, query, update, options);

        return result.getModifiedCount();
    }

    /**
     * Convert the given {@link Document} to an {@link EO}.
     */
    protected static EO docToObject(Document doc) {
        EO eo = new EO(doc.getString(KEY_TYPE), new EOBsonId(doc.getObjectId(KEY_ID)));
        Document dataDoc = (Document) doc.get(KEY_DATA);
        if (dataDoc != null) {
            for (String key : dataDoc.keySet()) {
                eo.setValue(key, dataDoc.get(key));
            }
        }

        return eo;
    }

    /**
     * Convert the given {@link Document} to an {@link EOAssociation}.
     */

    protected static EOAssociation docToAssoc(Document doc) {
        EOId srcId = new EOBsonId(doc.getObjectId(KEY_SRCOBJ));
        EOId dstId = new EOBsonId(doc.getObjectId(KEY_DSTOBJ));

        EOAssociation assoc = new EOAssociation(srcId, dstId, doc.getString(KEY_TYPE),
                Instant.ofEpochMilli(doc.getLong(KEY_CTIME)));

        Document dataDoc = (Document) doc.get(KEY_DATA);
        if (dataDoc != null) {
            for (String key : dataDoc.keySet()) {
                assoc.setValue(key, dataDoc.get(key));
            }
        }

        return assoc;
    }

    /**
     * Return all the associations of the given type for the given object. The cursor that is returned is tied
     * to the transaction; rolling back or committing the transaction and then consuming from the iterator will yield
     * undefined results.
     */
    public EOAssociationCursor assocGet(EOId id, String type) {
        return assocGet(id, type, null);
    }

    /**
     * Return all the associations of the given type for the given object. The cursor that is returned is tied
     * to the transaction; rolling back or committing the transaction and then consuming from the iterator will yield
     * undefined results.
     */
    public EOAssociationCursor assocGet(EOId id, String type, String name) {
        logger.debug("assocGet() [id: {}] [type: {}] [name: {}]", id, type, name);
        MongoCollection<Document> col = db.getCollection(COLLECTION_ASSOC);
        BasicDBObject query = new BasicDBObject();
        query.put(KEY_SRCOBJ, ((EOBsonId) id).getObjectId());
        query.put(KEY_TYPE, type);
        if (name != null) {
            query.put(KEY_DATA + "." + KEY_NAME, name);
        }

        logger.debug("Searching. [query: {}]", query);

        FindIterable<Document> result = col.find(session, query);
        MongoCursor<Document> it = result.iterator();
        return new EOAssociationCursor(it);
    }

    /**
     * Return the first association found in the cursor, and close the cursor.
     */
    public EOAssociation first(EOAssociationCursor cursor) {
        if (cursor != null) {
            Iterator<EOAssociation> it = cursor.iterator();
            if (!it.hasNext()) return null;
            EOAssociation assoc = it.next();
            cursor.close();
            return assoc;
        }
        return null;
    }


    /**
     * Return all the associations of the given type for the given object in descending order of ctime.
     * The iterable that is returned is tied to the transaction; rolling back or committing the transaction and
     * then consuming from the iterator will yield undefined results. Up to <code>limit</code> results are returned,
     * or specify <code>0</code> for no limit.
     */
    public EOAssociationCursor assocGetTimeRange(EOId id, String type, int limit) {
        MongoCollection<Document> col = db.getCollection(COLLECTION_ASSOC);
        BasicDBObject query = new BasicDBObject();
        query.put(KEY_SRCOBJ, ((EOBsonId) id).getObjectId());
        query.put(KEY_TYPE, type);

        FindIterable<Document> result = col.find(session, query).sort(new BasicDBObject(KEY_CTIME, -1)).limit(limit);
        MongoCursor<Document> it = result.iterator();
        return new EOAssociationCursor(it);
    }


    /**
     * Return all the matching objects of the given type. The order returned is undefined.
     */
    public EOObjectCursor objGet(String type) {
        MongoCollection<Document> col = db.getCollection(COLLECTION_OBJ);
        BasicDBObject query = new BasicDBObject();
        query.put(KEY_TYPE, type);

        FindIterable<Document> result = col.find(session, query);

        MongoCursor<Document> it = result.iterator();
        return new EOObjectCursor(it);
    }

    /**
     * Delete the given object. All associations which either originated from, or pointed to, this object are also
     * deleted.
     */
    public void objDelete(EOId id) {
        MongoCollection<Document> objCollection = db.getCollection(COLLECTION_OBJ);
        objCollection.deleteOne(session, new Document().append("_id", ((EOBsonId) id).getObjectId()));

        MongoCollection<Document> assocCollection = db.getCollection(COLLECTION_ASSOC);
        Document clause1 = new Document(KEY_SRCOBJ, ((EOBsonId) id).getObjectId());
        Document clause2 = new Document(KEY_DSTOBJ, ((EOBsonId) id).getObjectId());
        BasicDBList or = new BasicDBList();
        or.add(clause1);
        or.add(clause2);
        Document query = new Document("$or", or);
        assocCollection.deleteMany(session, query);
    }

    /**
     * Put a file into the database. Note that due to lack of support for multi-document transactions with GridFS,
     * the actual storage of the file is not transactional. If this transaction is rolled back an effort will be made
     * to delete the file; however the result of this operation is not guaranteed. If the file fails to delete for
     * whatever reason, the EOFS directory objects (created on the transaction) will be correct, but the file will
     * actually exist in GridFs, in an unknown condition; perhaps fully or partially intact.
     */
    public EOId fsFileAdd(EOId rootInodeId, String pathStr, InputStream file, String contentType) {
        Path path = Path.of(pathStr);
        Path filename = path.getFileName();
        Path parent = path.getParent();

        EOId parentInodeId = fsPathTraverse(rootInodeId, parent);
        logger.debug("Directory. [id: {}]", parentInodeId);

        if (parentInodeId == null) {
            logger.warn("Directory not found for file: {}", parent);
            return null;
        }

        // Step 1: Upload the file itself
        GridFSUploadOptions options = new GridFSUploadOptions()
                .chunkSizeBytes(1048576)
                .metadata(new Document(KEY_CONTENT_TYPE, contentType));
        GridFSBucket gridFSBucket = GridFSBuckets.create(db);
        // TODO: We don't really want to store the filename here, or do we?
        ObjectId gridFsId = gridFSBucket.uploadFromStream(filename.toString(), file, options);

        logger.debug("File updated. Gridid: {}", gridFsId);

        synchronized (this) {
            if (pendingFiles == null) pendingFiles = new HashSet<>();
            pendingFiles.add(gridFsId);
        }

        // Step 2: Create a file entry
        EOId fid = objAdd(TYPE_FS_FILE, Map.of(KEY_FS_FILENAME, pathStr, KEY_FS_GRIDFS_ID, gridFsId));
        assocAdd(parentInodeId, fid, TYPE_FS_LINK, Map.of(KEY_FS_FILENAME, filename.toString()));

        return fid;
    }

    /**
     * Find the file in the filesystem and write it to the provided {@link OutputStream}. The path specified
     * is relative to the inode specified by <code>rootInode</code>.
     */
    public void fsFileGet(EOId rootInode, OutputStream os, String pathStr) throws IOException {
        Path path = Path.of(pathStr);
        Path parent = path.getParent();
        Path filename = path.getFileName();

        logger.debug("Looking for directory: {}", parent);
        EOId inode = fsPathTraverse(rootInode, parent);

        logger.debug("Found directory: {}", inode);

        if (inode == null) {
            logger.warn("Path not found: {}", parent);
            // FIXME return or throw
            return;
        }

        EOAssociation link = first(assocGet(inode, TYPE_FS_LINK, filename.toString()));

        logger.debug("Searching for associations");
        try (EOAssociationCursor cursor = assocGet(inode, TYPE_FS_LINK)) {
            for (EOAssociation asc : cursor) {
                logger.debug("Asc: {}", asc);
            }
        }

        if (link == null) {
            logger.warn("File not found: {}", link);
            // FIXME return or throw
            return;
        }

        EO file = objGet(link.getDstObjectId());
        ObjectId gridfsId = (ObjectId) file.get(KEY_FS_GRIDFS_ID);
        logger.debug("Located gridfs id for file: {}", gridfsId);
        GridFSBucket gridFSBucket = GridFSBuckets.create(db);
        gridFSBucket.downloadToStream(gridfsId, os);
        os.flush();
    }

    /**
     * Create a new directory and link it to a parent. If a path consisting of multiple directories is
     * specified, all directories will be created (analogous to <code>mkdir -p</code> on UNIX systems).
     */
    public EOId fsDirAdd(EOId parentId, String path) {
        EOId bottomDirectory = null;
        EOId lastId = parentId;
        Path p = Path.of(path);

        // FIXME: Need to check if these are duplicate entries (or get an index?)
        for (Path part : p) {
            EOId newId = objAdd(TYPE_INODE);
            if (bottomDirectory == null) bottomDirectory = newId;
            logger.debug("Added path part: {}", part);
            assocAdd(lastId, newId, TYPE_FS_LINK, Map.of(KEY_FS_FILENAME, part.toString()));
            assocAdd(lastId, newId, TYPE_FS_LINK, Map.of(KEY_FS_FILENAME, VALUE_FS_PARENT));
            lastId = newId;
        }

        return bottomDirectory;
    }

    /**
     * Traverse the provided path relative to the provided inode.
     */
    protected EOId fsPathTraverse(EOId inode, Path p) {
        if (p.toString().equals(File.separator)) {
            return inode;
        }

        EOId cwd = inode;

        for (Path part : p) {
            logger.debug("Path part: {}", part);
            EOAssociation link = first(assocGet(cwd, TYPE_FS_LINK, part.toString()));
            logger.debug("Link: {}", link);
            if (link == null) {
                logger.warn("Directory not found: {}", part);
                return null;
            }
            cwd = link.getDstObjectId();
        }

        return cwd;
    }

    /**
     * Create a root directory inodes. This inodes will have some special properties; both <code>.</code> and
     * <code>..</code> will refer to the directory itself, in keeping with Unix tradition. There is no limit
     * to the number of root inodes which may be present within a database.
     */
    public EOId fsDirAdd() {
        EOId newId = objAdd(TYPE_INODE);
        assocAdd(newId, newId, TYPE_FS_LINK, Map.of(KEY_FS_FILENAME, VALUE_FS_PARENT));
        assocAdd(newId, newId, TYPE_FS_LINK, Map.of(KEY_FS_FILENAME, VALUE_FS_SELF));
        return newId;
    }

    /**
     * Commit the current transaction, and begin a new one.
     */
    public void commit() {
        session.commitTransaction();
        synchronized (this) {
            pendingFiles = null;
        }
        session.startTransaction();
    }

    /**
     * Roll back the current transaction and start a new one. If this transaction has pending files (files that
     * were created with <code>fileCreate()</code> but have not yet been comitted, then this method will make an
     * effort to delete those files from GridFS, which is not guaranteed. References to the file which were created
     * in the directory objects will be reliably rolled back.
     */
    public void rollback() {
        session.abortTransaction();
        synchronized (this) {
            if (pendingFiles != null) {
                GridFSBucket gridFSBucket = GridFSBuckets.create(db);
                for (ObjectId pending : pendingFiles) {
                    try {
                        gridFSBucket.delete(pending);
                    } catch (MongoGridFSException e) {
                        logger.warn("Failed to delete file created adjacent to pending transaction: {}", e);
                    }
                }
            }
            pendingFiles = null;
        }
        session.startTransaction();
    }

    /**
     * Close the given transaction and free any associated resources. If the transaction had pending changes, it will
     * be rolled back.
     */
    @Override
    public void close() {
        rollback();
        session.close();
        client.close();
    }

    /**
     * Drop the database associated with this transaction. Calling this method will automatically rollback the
     * transaction.
     */
    public void drop() {
        rollback();
        client.getDatabase(databaseName).drop();
    }
}
