package org.locationtech.geogig.model.internal;

import java.io.ByteArrayOutputStream;
import java.io.DataOutputStream;
import java.io.IOException;
import java.util.Map;
import java.util.NoSuchElementException;
import java.util.Set;
import java.util.SortedMap;
import java.util.TreeMap;

import org.rocksdb.ColumnFamilyDescriptor;
import org.rocksdb.ColumnFamilyHandle;
import org.rocksdb.ReadOptions;
import org.rocksdb.RocksDB;
import org.rocksdb.RocksDBException;
import org.rocksdb.WriteBatch;
import org.rocksdb.WriteOptions;

import com.google.common.base.Charsets;
import com.google.common.base.Preconditions;
import com.google.common.base.Throwables;
import com.google.common.collect.ImmutableSortedMap;
import com.google.common.io.ByteStreams;

class RocksdbNodeStore {

    private RocksDB db;

    private WriteOptions writeOptions;

    private ReadOptions readOptions;

    private ColumnFamilyHandle column;

    public RocksdbNodeStore(RocksDB db) {
        this.db = db;
        ColumnFamilyDescriptor columnDescriptor = new ColumnFamilyDescriptor(
                "nodes".getBytes(Charsets.UTF_8));
        try {
            column = db.createColumnFamily(columnDescriptor);
        } catch (RocksDBException e) {
            throw Throwables.propagate(e);
        }

        this.writeOptions = new WriteOptions();
        writeOptions.setDisableWAL(true);
        writeOptions.setSync(false);

        readOptions = new ReadOptions();
        readOptions.setFillCache(true).setVerifyChecksums(false);
    }

    public void close() {
        batch.close();
        readOptions.close();
        writeOptions.close();
        column.close();
        db = null;
    }

    public DAGNode get(NodeId nodeId) {
        flush();
        byte[] value;
        byte[] key = toKey(nodeId);
        try {
            value = db.get(column, readOptions, key);
            if (value == null) {
                throw new NoSuchElementException("Node " + nodeId + " not found");
            }
        } catch (Exception e) {
            throw Throwables.propagate(e);
        }
        return decode(value);
    }

    public SortedMap<NodeId, DAGNode> getAll(Set<NodeId> nodeIds) {
        if (nodeIds.isEmpty()) {
            return ImmutableSortedMap.of();
        }
        flush();
        SortedMap<NodeId, DAGNode> res = new TreeMap<>();
        byte[] valueBuff = new byte[512];
        nodeIds.forEach((id) -> {
            try {
                byte[] val = valueBuff;
                byte[] key = toKey(id);
                int ret = db.get(column, readOptions, key, val);
                Preconditions.checkState(ret != RocksDB.NOT_FOUND);
                if (ret > valueBuff.length) {
                    val = db.get(column, key);
                }
                DAGNode node = decode(val);
                res.put(id, node);
            } catch (IllegalArgumentException | RocksDBException e) {
                throw Throwables.propagate(e);
            }
        });
        return res;
    }

    private WriteBatch batch = new WriteBatch();

    private int batchSize;

    public void put(NodeId nodeId, DAGNode node) {
        batch.put(column, toKey(nodeId), encode(node));
        batchSize++;

        flush(1_000);
    }

    public void putAll(Map<NodeId, DAGNode> nodeMappings) {
        ByteArrayOutputStream out = new ByteArrayOutputStream();

        nodeMappings.forEach((nodeId, dagNode) -> {
            out.reset();
            encode(dagNode, out);
            byte[] value = out.toByteArray();
            byte[] key = toKey(nodeId);
            batch.put(column, key, value);
            batchSize++;
            flush(1_000);
        });
    }

    private void flush() {
        flush(1);
    }

    private void flush(int limit) {
        if (batchSize >= limit) {
            try {
                db.write(writeOptions, batch);
            } catch (Exception ex) {
                throw Throwables.propagate(ex);
            }
            batchSize = 0;
            batch.close();
            batch = new WriteBatch();
        }
    }

    private byte[] toKey(NodeId nodeId) {
        byte[] key = nodeId.name().getBytes(Charsets.UTF_8);
        return key;
    }

    private NodeId fromKey(byte[] key) {
        String name = new String(key, Charsets.UTF_8);
        return new CanonicalNodeId(name);
    }

    private byte[] encode(DAGNode node) {
        return encode(node, new ByteArrayOutputStream());
    }

    private byte[] encode(DAGNode node, ByteArrayOutputStream outstream) {
        DataOutputStream out = new DataOutputStream(outstream);
        try {
            DAGNode.encode(node, out);
        } catch (IOException e) {
            throw Throwables.propagate(e);
        }
        return outstream.toByteArray();
    }

    private DAGNode decode(byte[] nodeData) {
        DAGNode node;
        try {
            node = DAGNode.decode(ByteStreams.newDataInput(nodeData));
        } catch (IOException e) {
            throw Throwables.propagate(e);
        }
        return node;
    }
}

@@ -3131,32 +3131,67 @@
 DAGNode node) %7B%0A
+        synchronized (batch) %7B%0A    
         batch.pu
@@ -3230,32 +3230,36 @@
 node));%0A        
+    
 batchSize++;%0A%0A  
@@ -3254,16 +3254,26 @@
 hSize++;
+%0A        %7D
 %0A%0A      
@@ -3610,32 +3610,71 @@
  toKey(nodeId);%0A
+            synchronized (batch) %7B%0A    
             batc
@@ -3696,24 +3696,28 @@
 ey, value);%0A
+    
             
@@ -3721,32 +3721,46 @@
    batchSize++;%0A
+            %7D%0A
             flus
@@ -3902,32 +3902,71 @@
 ize %3E= limit) %7B%0A
+            synchronized (batch) %7B%0A    
             try 
@@ -3959,32 +3959,36 @@
           try %7B%0A
+    
                 
@@ -4018,16 +4018,20 @@
 batch);%0A
+    
         
@@ -4067,32 +4067,36 @@
                 
+    
 throw Throwables
@@ -4115,34 +4115,42 @@
 x);%0A            
-%7D%0A
+    %7D%0A    
             batc
@@ -4168,32 +4168,36 @@
             
+    
 batch.cl
 ose();%0A     
@@ -4176,35 +4176,35 @@
         batch.cl
-ose
+ear
 ();%0A            
@@ -4207,33 +4207,9 @@
     
-batch = new WriteBatch();
+%7D
 %0A   

Synchronize on RocksdbNodeStore.writeBatch