package main

import (
	"bytes"
	"reflect"
	"testing"

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

func TestTimeDecode(t *testing.T) {
	data := make(chan []byte, 1)
	out := make(chan combiner.Record, 1)

	var nanos int64 = 1000
	event := &pbmsg.Event{Time: proto.Int64(nanos)}

	buf, err := proto.Marshal(&pbmsg.EventSet{Event: []*pbmsg.Event{event}})
	if err != nil {
		t.Fatalf("proto.Marshal; err = %q", err)
	}

	data <- buf
	close(data)
	decode(out, data, keepAll)
	dec := <-out
	result := (*pbmsg.Event)(dec.(*eventRecord))

	if have, want := result, event; !reflect.DeepEqual(have, want) {
		t.Errorf("have = %v; want %v", have, want)
	}
}

var mockEvent = &pbmsg.Event{
	Kind:          pbmsg.Event_REQUEST_HEAD_RECEIVED.Enum(),
	Time:          proto.Int64(1446236540631000000),
	Hostname:      proto.String("trace.internal.justin.tv"),
	Svcname:       proto.String("code.justin.tv/release/trace/pbcollect"),
	Pid:           proto.Int32(1),
	TransactionId: []uint64{1, 2},
	Extra: &pbmsg.Extra{
		Http: &pbmsg.ExtraHTTP{
			StatusCode: proto.Int32(200),
			Method:     pbmsg.ExtraHTTP_GET.Enum(),
			UriPath:    proto.String("/notreal"),
		},
	},
}

var mockRecord *eventRecord = (*eventRecord)(mockEvent)

var mockEventSet = &pbmsg.EventSet{Event: []*pbmsg.Event{mockEvent}}

// Make an event set with n events
func mockEventSetLen(n int) *pbmsg.EventSet {
	events := make([]*pbmsg.Event, n)
	for i := range events {
		events[i] = mockEvent
	}
	return &pbmsg.EventSet{Event: events}
}

// Make a mock event of byte size n
func mockEventSetSize(t *testing.T, n int) *pbmsg.EventSet {
	// have to iteratively try to reach byte size N
	var es *pbmsg.EventSet
	marshaledSize := n + 1
	bytes := make([]byte, n+1)
	for marshaledSize > n {
		bytes = bytes[:len(bytes)-1]
		ev := &pbmsg.Event{
			TransactionId: []uint64{0},
			Hostname:      proto.String(string(bytes))}
		es = &pbmsg.EventSet{Event: []*pbmsg.Event{ev}}
		marshaledSize = proto.Size(es)
	}
	if marshaledSize < n {
		t.Fatalf("Unable to create an eventset with size exactly %d, best I can do is %d", n, marshaledSize)
	}
	return es
}

func TestConcurrentDecode(t *testing.T) {
	// Make an event blob.
	eventBytes, err := proto.Marshal(mockEventSet)
	if err != nil {
		t.Fatal("Unable to marshal mockEvent")
	}

	// Set up a decoder.
	nEvents := 1000
	data := make(chan []byte, nEvents)
	blobs := concurrentDecode(data)

	// Pump data into the decoer's input channel.
	go func() {
		for i := 0; i < nEvents; i++ {
			data <- eventBytes
		}
		close(data)
	}()

	// Count the number of messages received, and make sure they are valid.
	nReceived := 0
	done := make(chan struct{})
	go func() {
		defer close(done)
		for msg := range blobs {
			data := msg.Data
			es := &pbmsg.EventSet{}
			proto.Unmarshal(data, es)
			for _, ev := range es.Event {
				nReceived += 1
				if !reflect.DeepEqual(ev, mockEvent) {
					t.Errorf("Message changed while decoding into blob.")
					t.Logf("Have: %v", ev)
					t.Logf("Want: %v", mockEvent)
				}
			}
		}
	}()

	<-done
	if nReceived != nEvents {
		t.Errorf("Expected %d events, only found %d", nEvents, nReceived)
	}
}

// Test that the decode loop correctly unpacks its inputs.
func TestDecodeLoop(t *testing.T) {
	// Make two events with different transaction IDs
	e1 := mockEvent
	e2 := new(pbmsg.Event)
	*e2 = *e1
	e2.TransactionId = []uint64{e1.TransactionId[0] + 1, e1.TransactionId[1] + 1}

	if reflect.DeepEqual(e1.TransactionId, e2.TransactionId) {
		t.Fatal("should have different transaction IDs")
	}

	// Put them into two different eventsets
	es1 := &pbmsg.EventSet{Event: []*pbmsg.Event{e1}}
	es2 := &pbmsg.EventSet{Event: []*pbmsg.Event{e2}}

	// Put them into the data channel
	data := make(chan []byte, 2)
	data1, err := proto.Marshal(es1)
	if err != nil {
		t.Fatal(err)
	}
	data <- data1

	data2, err := proto.Marshal(es2)
	if err != nil {
		t.Fatal(err)
	}
	data <- data2
	close(data)

	// Do the decode
	out := make(chan combiner.Record, 2)
	decode(out, data, keepAll)
	close(out)

	// Pull out the output
	var have []combiner.Record
	for ev := range out {
		have = append(have, ev)
	}

	// It should be two eventRecords
	if len(have) != 2 {
		t.Fatalf("wrong number of events in output channel, have=%d  want=%d", len(have), 2)
	}

	ev1, ok := have[0].(*eventRecord)
	if !ok {
		t.Fatalf("unexpected type for record from decode: %T", have[0])
	}
	ev2, ok := have[1].(*eventRecord)
	if !ok {
		t.Fatalf("unexpected type for record from decode: %T", have[1])
	}

	// The first one should have the first input's transaction ID, the
	// second should have the second.
	if !reflect.DeepEqual(e1.TransactionId, ev1.TransactionId) {
		t.Errorf("first event transaction ID wrong")
	}
	if !reflect.DeepEqual(e2.TransactionId, ev2.TransactionId) {
		t.Errorf("second event transaction ID wrong")
	}
}

func TestDecoderSkipsInvalidEventSets(t *testing.T) {
	// Make an event blob.
	validBytes, err := proto.Marshal(mockEventSet)
	if err != nil {
		t.Fatal("Unable to marshal mockEventSet")
	}

	// Invalid bytes should be skipped
	invalidBytes := []byte{'h', 'e', 'l', 'l', 'o'}

	// Eventsets with multiple events should be skipped
	tooManyEvents, err := proto.Marshal(mockEventSetLen(2))
	if err != nil {
		t.Fatal("Unable to marshal tooManyEvents eventset")
	}

	tooBig, err := proto.Marshal(mockEventSetSize(t, MAX_RECORD_SIZE+1))
	if err != nil {
		t.Fatal("Unable to marshal tooBig eventSet")
	}

	// Set up a decoder.
	data := make(chan []byte, 7)
	blobs := concurrentDecode(data)

	data <- validBytes
	data <- invalidBytes

	data <- validBytes
	data <- tooManyEvents

	data <- validBytes
	data <- tooBig

	data <- validBytes
	close(data)

	// Count the number of messages received, and make sure they are valid.
	want := 4
	have := 0
	done := make(chan struct{})
	go func() {
		defer close(done)
		for msg := range blobs {
			data := msg.Data
			es := &pbmsg.EventSet{}
			proto.Unmarshal(data, es)
			for _, ev := range es.Event {
				have += 1
				if !reflect.DeepEqual(ev, mockEvent) {
					t.Errorf("Unexpected message out of blobs")
					t.Logf("Have: %v", ev)
					t.Logf("Want: %v", mockEvent)
				}
			}
		}
	}()

	<-done
	if have != want {
		t.Errorf("Expected %d events, only found %d", want, have)
	}

}

func TestEventRecordSetEncode(t *testing.T) {
	// map of description to pile of events to test encoding on
	tests := map[string][]*pbmsg.Event{
		"empty set":              []*pbmsg.Event{},
		"empty record":           []*pbmsg.Event{&pbmsg.Event{}},
		"single record":          []*pbmsg.Event{mockEvent},
		"two records":            []*pbmsg.Event{mockEvent, mockEvent},
		"empty record in middle": []*pbmsg.Event{mockEvent, &pbmsg.Event{}, mockEvent},
		"1000 records":           make([]*pbmsg.Event, 1000),
	}

	for i := range tests["1000 records"] {
		tests["1000 records"][i] = mockEvent
	}

	for desc, evs := range tests {
		realEventSet := &pbmsg.EventSet{Event: evs}
		want, err := proto.Marshal(realEventSet)
		if err != nil {
			t.Fatalf("marshal err test=%q err=%q", desc, err)
		}

		rc := newEventRecordSet()
		for _, ev := range evs {
			rc.Push((*eventRecord)(ev))
		}

		have := rc.Encode()
		if !bytes.Equal(have, want) {
			t.Errorf("failed encoding %s", desc)
		}
	}
}

func BenchmarkDecodeLoop(b *testing.B) {
	const (
		nIters = 100
	)
	msg, err := proto.Marshal(mockEventSet)
	if err != nil {
		b.Fatal(err)
	}

	out := make(chan combiner.Record, nIters)
	go func() {
		for range out {
		}
	}()

	b.SetBytes(int64(len(msg) * nIters))

	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		data := make(chan []byte, nIters)
		for i := 0; i < nIters; i++ {
			data <- msg
		}
		close(data)
		decode(out, data, keepAll)
	}
}

func BenchmarkEventRecordSetPush(b *testing.B) {
	rs := newEventRecordSet()
	b.ResetTimer()

	for i := 0; i < b.N; i++ {
		rs.Push(mockRecord)
	}
}

func BenchmarkEventRecordSetPushPop(b *testing.B) {
	rs := newEventRecordSet()
	for i := 0; i < 1000; i++ {
		rs.Push(mockRecord)
	}
	b.ResetTimer()

	for i := 0; i < b.N; i++ {
		rs.Push(mockRecord)
		rs.Pop()
	}
}

func BenchmarkEventRecordSetSize(b *testing.B) {
	rs := newEventRecordSet()
	for i := 0; i < 1000; i++ {
		rs.Push(mockRecord)
	}
	b.ResetTimer()

	for i := 0; i < b.N; i++ {
		rs.Size()
	}
}

func BenchmarkEventRecordSetEncode(b *testing.B) {
	rs := newEventRecordSet()
	for i := 0; i < 1000; i++ {
		rs.Push(mockRecord)
	}
	b.ResetTimer()

	for i := 0; i < b.N; i++ {
		rs.Encode()
	}
}
