package com.thinkaurelius.titan.diskstorage.locking.consistentkey;

import java.util.List;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import com.google.common.base.Preconditions;
import com.thinkaurelius.titan.diskstorage.StaticBuffer;
import com.thinkaurelius.titan.diskstorage.StorageException;
import com.thinkaurelius.titan.diskstorage.keycolumnvalue.Entry;
import com.thinkaurelius.titan.diskstorage.keycolumnvalue.KeyColumnValueStore;
import com.thinkaurelius.titan.diskstorage.keycolumnvalue.KeySliceQuery;
import com.thinkaurelius.titan.diskstorage.keycolumnvalue.StoreTransaction;
import com.thinkaurelius.titan.diskstorage.locking.Locker;
import com.thinkaurelius.titan.diskstorage.locking.PermanentLockingException;
import com.thinkaurelius.titan.diskstorage.util.KeyColumn;
import com.thinkaurelius.titan.diskstorage.util.RecordIterator;

/**
 * A {@link KeyColumnValueStore} wrapper intended for nontransactional stores
 * that forwards all <b>but</b> these two methods to an encapsulated store
 * instance:
 * 
 * <ul>
 *   <li>{@link #acquireLock(StaticBuffer, StaticBuffer, StaticBuffer, StoreTransaction)}</li>
 *   <li>{@link #mutate(StaticBuffer, List, List, StoreTransaction)}</li>
 * </ul>
 * 
 * This wrapper adds some logic to both of the overridden methods before calling
 * the encapsulated store's version.
 * <p>
 * This class, along with its collaborator class
 * {@link ExpectedValueCheckingTransaction}, track all {@code expectedValue}
 * arguments passed to {@code acquireLock} for each {@code StoreTransaction}.
 * When the transaction first {@code mutate(...)}s, the these classes cooperate
 * to check that all previously provided expected values match actual values,
 * throwing an exception and preventing mutation if a mismatch is detected.
 * <p>
 * This relies on a {@code Locker} instance supplied during construction for
 * locking.
 * 
 */
public class ExpectedValueCheckingStore implements KeyColumnValueStore {

    /**
     * Configuration setting key for the local lock mediator prefix
     */
    public static final String LOCAL_LOCK_MEDIATOR_PREFIX_KEY = "local-lock-mediator-prefix";
    
    private static final Logger log = LoggerFactory.getLogger(ExpectedValueCheckingStore.class);

    /**
     * Titan data store.
     */
    final KeyColumnValueStore dataStore;
    
    final Locker locker;
    
    public ExpectedValueCheckingStore(KeyColumnValueStore dataStore, Locker locker) {
        Preconditions.checkNotNull(dataStore);
        this.dataStore = dataStore;
        this.locker = locker;
    }

    public KeyColumnValueStore getDataStore() {
        return dataStore;
    }

    private StoreTransaction getBaseTx(StoreTransaction txh) {
        Preconditions.checkNotNull(txh);
        Preconditions.checkArgument(txh instanceof ExpectedValueCheckingTransaction);
        return ((ExpectedValueCheckingTransaction) txh).getBaseTransaction();
    }

    @Override
    public boolean containsKey(StaticBuffer key, StoreTransaction txh) throws StorageException {
        return dataStore.containsKey(key, getBaseTx(txh));
    }

    @Override
    public List<Entry> getSlice(KeySliceQuery query, StoreTransaction txh) throws StorageException {
        return dataStore.getSlice(query, getBaseTx(txh));
    }

    /**
     * {@inheritDoc}
     * 
     * <p/>
     * 
     * This implementation supports locking when {@code lockStore} is non-null.  
     */
    @Override
    public void mutate(StaticBuffer key, List<Entry> additions, List<StaticBuffer> deletions, StoreTransaction txh) throws StorageException {
        if (locker != null) {
            ExpectedValueCheckingTransaction tx = (ExpectedValueCheckingTransaction) txh;
            if (!tx.isMutationStarted()) {
                tx.mutationStarted();
                locker.checkLocks(tx.getConsistentTransaction());
                tx.checkExpectedValues();
            }
        }
        dataStore.mutate(key, additions, deletions, getBaseTx(txh));
    }
    
    /**
     * {@inheritDoc}
     * 
     * <p/>
     * 
     * This implementation supports locking when {@code lockStore} is non-null.
     * <p>
     * Consider the following scenario. This method is called twice with
     * identical key, column, and txh arguments, but with different
     * expectedValue arguments in each call. In testing, it seems titan's
     * graphdb requires that implementations discard the second expectedValue
     * and, when checking expectedValues vs actual values just prior to mutate,
     * only the initial expectedValue argument should be considered.
     */
    @Override
    public void acquireLock(StaticBuffer key, StaticBuffer column, StaticBuffer expectedValue, StoreTransaction txh) throws StorageException {
        if (locker != null) {
            ExpectedValueCheckingTransaction tx = (ExpectedValueCheckingTransaction) txh;
            if (tx.isMutationStarted())
                throw new PermanentLockingException("Attempted to obtain a lock after mutations had been persisted");
            KeyColumn lockID = new KeyColumn(key, column);
            log.debug("Attempting to acquireLock on {} ev={}", lockID, expectedValue);
            locker.writeLock(lockID, tx.getConsistentTransaction());
            tx.storeExpectedValue(this, lockID, expectedValue);
        } else {
            dataStore.acquireLock(key, column, expectedValue, getBaseTx(txh));
        }
    }

    @Override
    public RecordIterator<StaticBuffer> getKeys(StoreTransaction txh) throws StorageException {
        return dataStore.getKeys(getBaseTx(txh));
    }

    @Override
    public StaticBuffer[] getLocalKeyPartition() throws StorageException {
        return dataStore.getLocalKeyPartition();
    }

    @Override
    public String getName() {
        return dataStore.getName();
    }

    @Override
    public void close() throws StorageException {
        dataStore.close();
        // TODO close locker?
    }
    
    void deleteLocks(ExpectedValueCheckingTransaction tx) throws StorageException {
        locker.deleteLocks(tx.getConsistentTransaction());
    }
}

@@ -85,16 +85,77 @@
 .List;%0A%0A
+import com.thinkaurelius.titan.diskstorage.keycolumnvalue.*;%0A
 import o
@@ -167,24 +167,24 @@
 f4j.Logger;%0A
-
 import org.s
@@ -371,301 +371,8 @@
 on;%0A
-import com.thinkaurelius.titan.diskstorage.keycolumnvalue.Entry;%0Aimport com.thinkaurelius.titan.diskstorage.keycolumnvalue.KeyColumnValueStore;%0Aimport com.thinkaurelius.titan.diskstorage.keycolumnvalue.KeySliceQuery;%0Aimport com.thinkaurelius.titan.diskstorage.keycolumnvalue.StoreTransaction;%0A
 impo
@@ -3046,24 +3046,227 @@
 h));%0A    %7D%0A%0A
+    @Override%0A    public List%3CList%3CEntry%3E%3E getSlice(List%3CStaticBuffer%3E keys, SliceQuery query, StoreTransaction txh) throws StorageException %7B%0A        return dataStore.getSlice(keys, query, txh);%0A    %7D%0A%0A
     /**%0A    
@@ -5496,32 +5496,32 @@
 rageException %7B%0A
-
         return d
@@ -5553,32 +5553,353 @@
 x(txh));%0A    %7D%0A%0A
+    @Override%0A    public KeyIterator getKeys(KeyRangeQuery query, StoreTransaction txh) throws StorageException %7B%0A        return dataStore.getKeys(query, txh);%0A    %7D%0A%0A    @Override%0A    public KeyIterator getKeys(SliceQuery query, StoreTransaction txh) throws StorageException %7B%0A        return getKeys(query, txh);%0A    %7D%0A%0A
     @Override%0A  

add missing getKeys/getSlice methods to ExpectedValueCheckingStore