/**
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 * http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package com.lmax.collections.coalescing.ring.buffer;

import java.util.Collection;
import java.util.concurrent.atomic.AtomicLong;
import java.util.concurrent.atomic.AtomicReferenceArray;

import static java.lang.Math.min;

public final class CoalescingRingBuffer<K, V> implements CoalescingBuffer<K, V> {

    private volatile long nextWrite = 1; // the next write index
    private long lastCleaned = 0; // the last index that was nulled out by the producer
    private final AtomicLong rejectionCount = new AtomicLong(0);
    private final K[] keys;
    private final AtomicReferenceArray<V> values;

    @SuppressWarnings("unchecked")
    private final K nonCollapsibleKey = (K) new Object();
    private final int mask;
    private final int capacity;

    private volatile long nextRead = 1; // the oldest slot that is is safe to write to
    private final AtomicLong lastRead = new AtomicLong(0); // the newest slot that it is safe to overwrite

    @SuppressWarnings("unchecked")
    public CoalescingRingBuffer(int capacity) {
        this.capacity = nextPowerOfTwo(capacity);
        this.mask = this.capacity - 1;


        this.keys = (K[]) new Object[this.capacity];
        this.values = new AtomicReferenceArray<V>(this.capacity);
    }

    private int nextPowerOfTwo(int value) {
        return 1 << (32 - Integer.numberOfLeadingZeros(value - 1));
    }

    @Override
    public int size() {
        return (int) (nextWrite - lastRead.get() - 1);
    }

    @Override
    public int capacity() {
        return capacity;
    }

    public long rejectionCount() {
        return rejectionCount.get();
    }

    public long nextWrite() {
        return nextWrite;
    }

    public long nextRead() {
        return nextRead;
    }

    @Override
    public boolean isEmpty() {
        return nextRead == nextWrite;
    }

    @Override
    public boolean isFull() {
        return size() == capacity;
    }

    @Override
    public boolean offer(K key, V value) {
        long nextWrite = this.nextWrite;

        for (long readPosition = nextRead; readPosition < nextWrite; readPosition++) {
            int index = mask(readPosition);

            if(key.equals(keys[index])) {
                values.set(index, value);

                if (nextRead <= readPosition) {  // check that the reader has not read it yet
                    return true;
                } else {
                    break;
                }
            }
        }

        return add(key, value);
    }

    @Override
    public boolean offer(V value) {
        return add(nonCollapsibleKey, value);
    }

    private boolean add(K key, V value) {
        if (isFull()) {
            rejectionCount.lazySet(rejectionCount.get() + 1);
            return false;
        }

        cleanUp();
        store(key, value);
        return true;
    }

    private void cleanUp() {
        long lastRead = this.lastRead.get();

        if (lastRead == lastCleaned) {
            return;
        }

        while (lastCleaned < lastRead) {
            int index = mask(++lastCleaned);
            keys[index] = null;
            values.lazySet(index, null);
        }
    }

    private void store(K key, V value) {
        long nextWrite = this.nextWrite;
        int index = mask(nextWrite);

        keys[index] = key;
        values.set(index, value);

        this.nextWrite = nextWrite + 1;
    }

    @Override
    public int poll(Collection<? super V> bucket) {
        claimUpTo(nextWrite);
        return fill(bucket);
    }

    @Override
    public int poll(Collection<? super V> bucket, int maxItems) {
        claimUpTo(min(nextRead + maxItems, nextWrite));
        return fill(bucket);
    }

    private void claimUpTo(long claimIndex) {
        nextRead = claimIndex;
    }

    private int fill(Collection<? super V> bucket) {
        long nextRead = this.nextRead;
        long lastRead = this.lastRead.get();

        for (long readIndex = lastRead + 1; readIndex < nextRead; readIndex++) {
            int index = mask(readIndex);
            bucket.add(values.get(index));
        }

        int readCount = (int) (nextRead - lastRead - 1);
        this.lastRead.lazySet(nextRead - 1);
        return readCount;
    }

    private int mask(long value) {
        return ((int) value) & mask;
    }

}
@@ -4016,38 +4016,8 @@
 ) %7B%0A
-        claimUpTo(nextWrite);%0A
     
@@ -4030,32 +4030,43 @@
 turn fill(bucket
+, nextWrite
 );%0A    %7D%0A%0A    @O
@@ -4147,16 +4147,21 @@
         
+long 
 claimUpT
@@ -4161,17 +4161,19 @@
 laimUpTo
-(
+ = 
 min(next
@@ -4199,17 +4199,16 @@
 xtWrite)
-)
 ;%0A      
@@ -4231,93 +4231,20 @@
 cket
-);%0A    %7D%0A%0A    private void claimUpTo(long claimIndex) %7B%0A        nextRead = claimIndex
+, claimUpTo)
 ;%0A  
@@ -4289,32 +4289,48 @@
  super V%3E bucket
+, long claimUpTo
 ) %7B%0A        long
@@ -4321,29 +4321,24 @@
 ) %7B%0A        
-long 
 nextRead = t
@@ -4340,21 +4340,17 @@
 d = 
-this.nextRead
+claimUpTo
 ;%0A  
@@ -4449,24 +4449,25 @@
 Index %3C 
-nextRead
+claimUpTo
 ; readIn
@@ -4574,65 +4574,8 @@
  %7D%0A%0A
-        int readCount = (int) (nextRead - lastRead - 1);%0A
     
@@ -4600,24 +4600,25 @@
 lazySet(
-nextRead
+claimUpTo
  - 1);%0A 
@@ -4635,17 +4635,40 @@
 urn 
-readCount
+(int) (claimUpTo - lastRead - 1)
 ;%0A  

remove unnecessary volatile read during poll