package sink

import (
	"encoding/binary"
	"expvar"
	"fmt"
	"hash"
	"math"
	"sync"
	"time"

	"crypto/sha1"

	"code.justin.tv/ear/kinesis/combiner"
)

// SplitSink allocates blobs between two sinks. With probability p, a
// blob will go to s1; with probabilty (1-p), it will go to s2, and it
// will never go to both. Sharding is based on the blobs' keys.
//
// SplitSink uses non-blockings ends to write to each sink. If a
// sink's buffer fills up, it will skip sending blobs to the sink;
// they won't go to the other sink, they'll just be dropped
// entirely. The number of skipped and accepted blobs are recorded in
// the package's expvar data under the key 'SplitSink'; if data is
// already stored under this key (like from a previous call of
// SplitSink) then it will be overwritten.
func SplitSink(p float64, s1, s2 Sink) Sink {
	return newSplitSink(p, s1, s2).run
}

// how long a sink has to accept a record before it will be skipped
const splitSinkReceiveTimeout = time.Millisecond

type splitSink struct {
	h      hash.Hash
	s1, s2 Sink
	p      float64
}

func newSplitSink(p float64, s1, s2 Sink) *splitSink {
	return &splitSink{
		h:  sha1.New(),
		s1: s1,
		s2: s2,
		p:  p,
	}
}

func (s *splitSink) run(data <-chan *combiner.Blob) {
	var wg sync.WaitGroup
	c1 := make(chan *combiner.Blob, cap(data))
	c2 := make(chan *combiner.Blob, cap(data))

	statGroup := new(expvar.Map).Init()
	expStats.Set("SplitSink", statGroup)
	stat1 := newSinkStat(c1)
	stat2 := newSinkStat(c2)
	statGroup.Set(fmt.Sprintf("%s-1", getFunctionName(s.s1)), stat1)
	statGroup.Set(fmt.Sprintf("%s-2", getFunctionName(s.s2)), stat2)

	wg.Add(2)
	go func() {
		s.s1(c1)
		wg.Done()
	}()

	go func() {
		s.s2(c2)
		wg.Done()
	}()

	var (
		target chan *combiner.Blob
		stat   *sinkStat
	)

	for b := range data {
		if s.rand(b) < s.p {
			target = c1
			stat = stat1
		} else {
			target = c2
			stat = stat2
		}
		// Do a nonblocking send to the receiver. If it fails, set a
		// timer; doing this under the default case avoids allocating
		// a bajillion timers except when buffers are full. It also
		// gives the sink a chance to execute if the SplitSink
		// suddenly gets a rush of data.
		select {
		case target <- b:
			// The happy case: The buffer isn't full.
			stat.RecordAccepted()
		default:
			// Unhappy case: buffer is full. Set a timeout to give
			// the receiver a chance to catch its breath.
			t := time.NewTimer(splitSinkReceiveTimeout)
			select {
			case target <- b:
				stat.RecordAccepted()
			case <-t.C:
				stat.RecordSkipped()
			}
			t.Stop()
		}
	}

	close(c1)
	close(c2)

	wg.Wait()

}

// Generates a pseudorandom number in [0.0, 1.0] from a blob's key
func (s *splitSink) rand(b *combiner.Blob) float64 {
	s.h.Reset()
	s.h.Write([]byte(b.Key))
	val := binary.LittleEndian.Uint64(s.h.Sum(nil))
	return float64(val) / float64(math.MaxUint64)
}
