package com.facebook.presto.operator.tap;

import com.facebook.presto.Tuple;
import com.facebook.presto.block.TupleStreamPosition;
import com.facebook.presto.slice.Slice;
import com.facebook.presto.slice.SliceInput;
import com.facebook.presto.slice.SliceOutput;

import static com.google.common.base.Preconditions.checkNotNull;
import static com.google.common.base.Preconditions.checkState;

public class StatsTupleValueSink
        implements TupleValueSink
{
    private long rowCount;
    private long runsCount;
    private Tuple lastTuple;
    private long minPosition = Long.MAX_VALUE;
    private long maxPosition = -1;
    private boolean finished = false;

    @Override
    public void process(TupleStreamPosition tupleStreamPosition)
    {
        checkNotNull(tupleStreamPosition, "tupleStreamPosition is null");
        checkState(!finished, "already finished");
        if (lastTuple == null) {
            lastTuple = tupleStreamPosition.getTuple();
        }
        else if (!tupleStreamPosition.currentTupleEquals(lastTuple)) {
            runsCount++;
            lastTuple = tupleStreamPosition.getTuple();
        }
        minPosition = Math.min(minPosition, tupleStreamPosition.getPosition());
        maxPosition = Math.max(maxPosition, tupleStreamPosition.getCurrentValueEndPosition());
        rowCount += tupleStreamPosition.getCurrentValueEndPosition() - tupleStreamPosition.getPosition() + 1;
    }

    @Override
    public void finished()
    {
        finished = true;
    }

    public Stats getStats()
    {
        return new Stats(rowCount, runsCount + 1, minPosition, maxPosition, rowCount / (runsCount + 1));
    }

    public static class Stats
    {
        private final long rowCount;
        private final long runsCount;
        private final long minPosition;
        private final long maxPosition;
        private final long avgRunLength;

        public Stats(long rowCount, long runsCount, long minPosition, long maxPosition, long avgRunLength)
        {
            this.rowCount = rowCount;
            this.runsCount = runsCount;
            this.minPosition = minPosition;
            this.maxPosition = maxPosition;
            this.avgRunLength = avgRunLength;
        }
        
        public static void serialize(Stats stats, SliceOutput sliceOutput)
        {
            // TODO: add a better way of serializing the stats that is less fragile
            sliceOutput.appendLong(stats.getRowCount())
                    .appendLong(stats.getRunsCount())
                    .appendLong(stats.getMinPosition())
                    .appendLong(stats.getMaxPosition())
                    .appendLong(stats.getAvgRunLength());
        }
        
        public static Stats deserialize(Slice slice)
        {
            SliceInput input = slice.input();
            long rowCount = input.readLong();
            long runsCount = input.readLong();
            long minPosition = input.readLong();
            long maxPosition = input.readLong();
            long avgRunLength = input.readLong();
            return new Stats(rowCount, runsCount, minPosition, maxPosition, avgRunLength);
        }

        public long getRowCount()
        {
            return rowCount;
        }

        public long getRunsCount()
        {
            return runsCount;
        }

        public long getMinPosition()
        {
            return minPosition;
        }

        public long getMaxPosition()
        {
            return maxPosition;
        }

        public long getAvgRunLength()
        {
            return avgRunLength;
        }
    }
}

@@ -256,16 +256,65 @@
 utput;%0A%0A
+import java.util.HashSet;%0Aimport java.util.Set;%0A%0A
 import s
@@ -503,16 +503,71 @@
 eSink%0A%7B%0A
+    private static final int MAX_UNIQUE_COUNT = 1000;%0A%0A
     priv
@@ -724,16 +724,84 @@
 n = -1;%0A
+    private final Set%3CTuple%3E set = new HashSet%3C%3E(MAX_UNIQUE_COUNT);%0A
     priv
@@ -1122,32 +1122,131 @@
 ion.getTuple();%0A
+            if (set.size() %3C MAX_UNIQUE_COUNT) %7B%0A                set.add(lastTuple);%0A            %7D%0A
         %7D%0A      
@@ -1383,32 +1383,131 @@
 ion.getTuple();%0A
+            if (set.size() %3C MAX_UNIQUE_COUNT) %7B%0A                set.add(lastTuple);%0A            %7D%0A
         %7D%0A      
@@ -2007,16 +2007,83 @@
 unt + 1)
+, (set.size() == MAX_UNIQUE_COUNT) ? Integer.MAX_VALUE : set.size()
 );%0A    %7D
@@ -2315,16 +2315,55 @@
 nLength;
+%0A        private final int uniqueCount;
 %0A%0A      
@@ -2461,16 +2461,33 @@
 unLength
+, int uniqueCount
 )%0A      
@@ -2694,32 +2694,76 @@
 = avgRunLength;%0A
+            this.uniqueCount = uniqueCount;%0A
         %7D%0A      
@@ -2752,32 +2752,24 @@
 ;%0A        %7D%0A
-        
 %0A        pub
@@ -3204,16 +3204,71 @@
 ength())
+%0A                    .appendInt(stats.getUniqueCount())
 ;%0A      
@@ -3271,24 +3271,16 @@
       %7D%0A
-        
 %0A       
@@ -3614,32 +3614,79 @@
 put.readLong();%0A
+            int uniqueCount = input.readInt();%0A
             retu
@@ -3757,16 +3757,29 @@
 unLength
+, uniqueCount
 );%0A     
@@ -4220,21 +4220,110 @@
 h;%0A        %7D
+%0A%0A        public int getUniqueCount()%0A        %7B%0A            return uniqueCount;%0A        %7D
 %0A    %7D%0A%7D%0A

Add simplistic unique count to stats collector