package com.netflix.astyanax.recipes.uniqueness;

import java.util.List;

import com.google.common.base.Supplier;
import com.google.common.collect.Lists;
import com.netflix.astyanax.Keyspace;
import com.netflix.astyanax.MutationBatch;
import com.netflix.astyanax.connectionpool.exceptions.ConnectionException;
import com.netflix.astyanax.connectionpool.exceptions.NotFoundException;
import com.netflix.astyanax.model.Column;
import com.netflix.astyanax.model.ColumnFamily;
import com.netflix.astyanax.model.ColumnList;
import com.netflix.astyanax.model.ConsistencyLevel;
import com.netflix.astyanax.recipes.locks.BusyLockException;
import com.netflix.astyanax.recipes.locks.StaleLockException;

/**
 * Multi-row uniqueness constraint where all involved column families are dedicated
 * only to uniquness constraint.
 * @author elandau
 *
 * @param <C>
 */
public class DedicatedMultiRowUniquenessConstraint<C> implements UniquenessConstraint {

    private final Keyspace keyspace;

    private Integer ttl = null;
    private ConsistencyLevel consistencyLevel = ConsistencyLevel.CL_LOCAL_QUORUM;
    private final C uniqueColumnName;
    
    private class Row<K> {
        private final ColumnFamily<K, C> columnFamily;
        private final K row;
        
        Row(ColumnFamily<K, C> columnFamily, K row) {
            super();
            this.columnFamily = columnFamily;
            this.row = row;
        }
        
        void fillMutation(MutationBatch m, Integer ttl) {
            m.withRow(columnFamily, row).putEmptyColumn(uniqueColumnName, ttl);
        }
        
        void fillReleaseMutation(MutationBatch m) {
            m.withRow(columnFamily, row).deleteColumn(uniqueColumnName);
        }
        
        void verifyLock() throws Exception {
            // Phase 2: Read back all columns. There should be only 1
            ColumnList<C> result = keyspace.prepareQuery(columnFamily).setConsistencyLevel(consistencyLevel)
                    .getKey(row).execute().getResult();

            if (result.size() != 1) {
                throw new NotUniqueException(row.toString());
            }
        }
        
        Column<C> getUniqueColumn() throws ConnectionException {
            ColumnList<C> columns = keyspace.prepareQuery(columnFamily).getKey(row).execute().getResult();
            Column<C> foundColumn = null;
            for (Column<C> column : columns) {
                if (column.getTtl() == 0) {
                    if (foundColumn != null)
                        throw new RuntimeException("Row has duplicate quite columns");
                    foundColumn = column;
                }
            }
            if (foundColumn == null) {
                throw new NotFoundException("Unique column not found");
            }
            return foundColumn;
        }
    }
    
    private final List<Row<?>> locks = Lists.newArrayList();

    public DedicatedMultiRowUniquenessConstraint(Keyspace keyspace, Supplier<C> uniqueColumnSupplier) {
        this.keyspace = keyspace;
        this.uniqueColumnName = uniqueColumnSupplier.get();
    }

    public DedicatedMultiRowUniquenessConstraint(Keyspace keyspace, C uniqueColumnName) {
        this.keyspace = keyspace;
        this.uniqueColumnName = uniqueColumnName;
    }

    /**
     * TTL to use for the uniquness operation. This is the TTL for the columns
     * to expire in the event of a client crash before the uniqueness can be
     * committed
     * 
     * @param ttl
     * @return
     */
    public DedicatedMultiRowUniquenessConstraint<C> withTtl(Integer ttl) {
        this.ttl = ttl;
        return this;
    }

    /**
     * Consistency level used
     * 
     * @param consistencyLevel
     * @return
     */
    public DedicatedMultiRowUniquenessConstraint<C> withConsistencyLevel(ConsistencyLevel consistencyLevel) {
        this.consistencyLevel = consistencyLevel;
        return this;
    }

    /**
     * Add a row to the set of rows being tested for uniqueness
     * 
     * @param columnFamily
     * @param rowKey
     * @return
     */
    public <K> DedicatedMultiRowUniquenessConstraint<C> withRow(ColumnFamily<K, C> columnFamily, K rowKey) {
        locks.add(new Row<K>(columnFamily, rowKey));
        return this;
    }
    
    public C getLockColumn() {
        return uniqueColumnName;
    }

    @Override
    public void acquire() throws NotUniqueException, Exception {
        acquireAndMutate(null);
    }
    
    @Override
    public void acquireAndMutate(MutationBatch other) throws NotUniqueException, Exception {
        // Insert lock check column for all rows in a single batch mutation
        try {
            MutationBatch m = keyspace.prepareMutationBatch().setConsistencyLevel(consistencyLevel);
            for (Row<?> lock : locks) {
                lock.fillMutation(m, ttl);
            }
            m.execute();

            // Check each lock in order
            for (Row<?> lock : locks) {
                lock.verifyLock();
            }

            // Commit the unique columns
            for (Row<?> lock : locks) {
                lock.fillMutation(m, null);
            }
            
            if (other != null) {
                m.mergeShallow(other);
            }
            m.execute();
        }
        catch (BusyLockException e) {
            release();
            throw new NotUniqueException(e);
        }
        catch (StaleLockException e) {
            release();
            throw new NotUniqueException(e);
        }
        catch (Exception e) {
            release();
            throw e;
        }
    }

    @Override
    public void release() throws Exception {
        MutationBatch m = keyspace.prepareMutationBatch();
        for (Row<?> lock : locks) {
            lock.fillReleaseMutation(m);
        }
        m.execute();
    }
    
    /**
     * @return
     * @throws Exception
     */
    public Column<C> getUniqueColumn() throws Exception {
        if (locks.size() == 0) 
            throw new IllegalStateException("Missing call to withRow to add rows to the uniqueness constraint");
        
        // Get the unique row from all columns
        List<Column<C>> columns = Lists.newArrayList();
        for (Row<?> row : locks) {
            columns.add(row.getUniqueColumn());
        }
        
        // Check that all rows have the same unique column otherwise they are not part of 
        // the same unique group
        Column<C> foundColumn = columns.get(0);
        for (int i = 1; i < columns.size(); i++) {
            Column<C> nextColumn = columns.get(i);
            if (!nextColumn.getRawName().equals(foundColumn.getRawName())) {
                throw new NotUniqueException("The provided rows are not part of the same uniquness constraint");
            }
            
            if (foundColumn.hasValue() != nextColumn.hasValue()) {
                throw new NotUniqueException("The provided rows are not part of the same uniquness constraint");
            }
            
            if (foundColumn.hasValue() && !nextColumn.getByteBufferValue().equals((foundColumn.getByteBufferValue()))) {
                throw new NotUniqueException("The provided rows are not part of the same uniquness constraint");
            }
        }
        
        return foundColumn;
    }
}
