package main

import (
	"log"
	"sync"
	"time"

	"code.justin.tv/ear/kinesis/combiner"
	"code.justin.tv/release/trace/pbmsg"
	"github.com/golang/protobuf/proto"
)

const (
	// Number of goroutines decoding data
	decodeWorkers = 64

	// Maximum size of a single record output by the combiner
	MAX_RECORD_SIZE = 100000 // 100KB

	// Frequency to flush records from the combiner if it doesnt hit
	// MAX_RECORD_SIZE
	combinerFlushInterval = 10 * time.Second
)

// decode converts the stream of bytes into decoded messages, putting
// them in the out channel, and may sample on transaction ID.
func decode(out chan<- combiner.Record, data <-chan []byte, sampler func([2]uint64) bool) {
	pb := proto.NewBuffer(nil)
	var in pbmsg.EventSet
	for buf := range data {
		// We expect each []byte to contain exactly one pbmsg.EventSet, since
		// they come from newScanner.

		pb.SetBuf(buf)
		in.Reset()

		err := pb.Unmarshal(&in)
		if err != nil {
			stats.Inc("errors.unmarshal", 1, 1)
			continue
		}

		if len(in.Event) != 1 {
			stats.Inc("errors.eventset_size", 1, 1)
			continue
		}

		var txid [2]uint64
		copy(txid[:], in.Event[0].TransactionId)
		if !sampler(txid) {
			continue
		}
		stats.Inc("events_received", 1, 0.1) // Sample at 10% since this is a hot stat
		out <- (*eventRecord)(in.Event[0])
	}
}

// decode from data to out using decodeWorkers goroutines
func concurrentDecode(data <-chan []byte) <-chan *combiner.Blob {
	records := make(chan combiner.Record, cap(data))
	blobber := combiner.NewRecordCombiner(virtualShards, MAX_RECORD_SIZE, combinerFlushInterval)

	var wg sync.WaitGroup
	for i := 0; i < decodeWorkers; i++ {
		wg.Add(1)
		go func() {
			defer wg.Done()
			decode(records, data, defaultSample)
		}()
	}
	go func() {
		wg.Wait()
		close(records)
	}()

	go func() {
		blobber.Consume(records)
		for err := range blobber.Errors() {
			log.Printf("record combiner err=%q", err)
			stats.Inc("errors.combiner", 1, 1)
		}
	}()

	return blobber.Blobs()
}

type eventRecord pbmsg.Event

func (e *eventRecord) Key() int {
	return int(e.TransactionId[0])
}

func (e *eventRecord) NewSet() combiner.RecordSet {
	return newEventRecordSet()
}

type eventRecordSet struct {
	events *pbmsg.EventSet
	buf    *proto.Buffer
}

func newEventRecordSet() *eventRecordSet {
	return &eventRecordSet{
		events: &pbmsg.EventSet{
			Event: make([]*pbmsg.Event, 0),
		},
		buf: proto.NewBuffer(nil),
	}
}

func (es *eventRecordSet) Push(r combiner.Record) {

	es.events.Event = append(es.events.Event, (*pbmsg.Event)(r.(*eventRecord)))
	// Append the event to the data, as well. Eventsets can be simply
	// concatenated.
	es.buf.Marshal(&pbmsg.EventSet{Event: []*pbmsg.Event{(*pbmsg.Event)(r.(*eventRecord))}})
}

func (es *eventRecordSet) Pop() combiner.Record {
	var r *pbmsg.Event
	r, es.events.Event = es.events.Event[len(es.events.Event)-1], es.events.Event[:len(es.events.Event)-1]
	es.buf.Reset()
	es.buf.Marshal(es.events)
	return (*eventRecord)(r)
}

func (es *eventRecordSet) Size() int {
	return len(es.buf.Bytes())
}

func (es *eventRecordSet) Encode() []byte {
	return es.buf.Bytes()
}
