// package kinesalite provides helpers for using Kinesalite in Trace's tests
package kinesalite

import (
	"errors"
	"testing"
	"time"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/awserr"
	"github.com/aws/aws-sdk-go/aws/credentials"
	"github.com/aws/aws-sdk-go/aws/request"
	"github.com/aws/aws-sdk-go/aws/session"
	"github.com/aws/aws-sdk-go/service/kinesis"
)

const (
	// Address where a kinesalite instance is running
	addr = "http://127.0.0.1:4567"
	// How long to sleep after receiving response from kinesalite to
	// let it complete callbacks. This is required because Kinesalite
	// creates objects like ShardIterators in callbacks _after_ it has
	// responded to requests. We need to wait for it to do all of its
	// tasks.
	cbSleep = 2 * time.Millisecond
)

func Setup(t *testing.T, stream string, shards int) *kinesis.Kinesis {
	k := TestClient()
	err := CreateStream(k, stream, shards)
	if err != nil {
		t.Fatalf("createStream err=%q", err)
	}
	return k
}

func Teardown(t *testing.T, stream string) {
	k := TestClient()
	err := DeleteStream(k, stream)
	if err != nil {
		t.Fatalf("deleteStream err=%q", err)
	}
}

// AssertStreamContains verifies that all []bytes in want are present in the
// given topic. Doesn't enforce ordering of the byte slices, since
// that's not really one of Kinesis's guarantees.
func AssertStreamContains(t *testing.T, k *kinesis.Kinesis, stream string, want [][]byte) {
	recs, err := getAllRecords(k, stream)
	if err != nil {
		t.Fatalf("error getting records=%q", err.Error())
	}

	have := make(map[string]int)
	for _, rec := range recs {
		have[string(rec.Data)] += 1
	}

	for _, blob := range want {
		numExpected, ok := have[string(blob)]
		if !ok {
			t.Errorf("can't find %q in stream", string(blob))
			continue
		}
		if numExpected < 1 {
			t.Errorf("found %q in stream more often than expected", string(blob))
			continue
		}
		have[string(blob)] -= 1
	}
	for blob, count := range have {
		if count > 0 {
			t.Errorf("found %q in stream, wasn't expecting it", blob)
		}
	}
}

func getAllRecords(k *kinesis.Kinesis, stream string) (recs []*kinesis.Record, err error) {
	shards, err := listShards(k, stream)
	if err != nil {
		return nil, err
	}

	for _, shard := range shards {
		var (
			err       error
			shardRecs []*kinesis.Record
		)
		for i := 0; i < 5; i++ {
			shardRecs, err = getShardRecords(k, shard, stream)
			if err == nil {
				break
			}
		}
		if err != nil {
			return nil, err
		}
		recs = append(recs, shardRecs...)
	}
	return recs, nil
}

func getShardRecords(k *kinesis.Kinesis, shard, stream string) (recs []*kinesis.Record, err error) {
	itRequest := &kinesis.GetShardIteratorInput{
		ShardId:           aws.String(shard),
		StreamName:        aws.String(stream),
		ShardIteratorType: aws.String(kinesis.ShardIteratorTypeTrimHorizon),
	}
	it, err := k.GetShardIterator(itRequest)
	if err != nil {
		return nil, err
	}
	recRequest := &kinesis.GetRecordsInput{
		Limit:         aws.Int64(10000),
		ShardIterator: aws.String(aws.StringValue(it.ShardIterator)),
	}
	for {
		resp, err := k.GetRecords(recRequest)
		if err != nil {
			return nil, err
		}
		recs = append(recs, resp.Records...)
		// If the NextShardIterator is null, then the shard has been
		// closed, so we can't keep reading.
		if resp.NextShardIterator == nil {
			break
		}
		// If we're matched up with the head of the topic, then there
		// are no new records, so stop reading.
		if aws.Int64Value(resp.MillisBehindLatest) == 0 {
			break
		}
		recRequest.ShardIterator = aws.String(aws.StringValue(resp.NextShardIterator))
	}
	return recs, nil
}

func listShards(k *kinesis.Kinesis, name string) (ids []string, err error) {
	req := &kinesis.DescribeStreamInput{
		StreamName: aws.String(name),
		Limit:      aws.Int64(10000),
	}

	resp, err := k.DescribeStream(req)
	if err != nil {
		return nil, err
	}
	if resp.StreamDescription == nil {
		return nil, errors.New("empty DescribeStream response")
	}
	ids = make([]string, len(resp.StreamDescription.Shards))
	for i, shard := range resp.StreamDescription.Shards {
		ids[i] = aws.StringValue(shard.ShardId)
	}

	// If the response indicates pagination is necessary, read
	// through the pages
	for i := 0; i < 10 && aws.BoolValue(resp.StreamDescription.HasMoreShards); i++ {
		req.ExclusiveStartShardId = aws.String(ids[len(ids)-1])
		resp, err = k.DescribeStream(req)
		if err != nil {
			return nil, err
		}
		if resp.StreamDescription == nil {
			return nil, errors.New("empty DescribeStream response")
		}

		for _, shard := range resp.StreamDescription.Shards {
			ids = append(ids, aws.StringValue(shard.ShardId))
		}
	}

	return ids, nil
}

func CreateStream(k *kinesis.Kinesis, name string, shards int) error {
	create := &kinesis.CreateStreamInput{
		StreamName: aws.String(name),
		ShardCount: aws.Int64(int64(shards)),
	}

	if _, err := k.CreateStream(create); err != nil {
		return err
	}

	describe := &kinesis.DescribeStreamInput{
		StreamName: aws.String(name),
	}
	// poll until the stream is actually created
	for i := 1; i < 1000; i++ {
		resp, err := k.DescribeStream(describe)
		if err != nil {
			return err
		}
		if aws.StringValue(resp.StreamDescription.StreamStatus) == kinesis.StreamStatusActive {
			return nil
		} else {
			time.Sleep(1 * time.Millisecond)
		}
	}
	return errors.New("timed out waiting for stream to appear")
}

func DeleteStream(k *kinesis.Kinesis, name string) error {
	delete := &kinesis.DeleteStreamInput{
		StreamName: aws.String(name),
	}
	if _, err := k.DeleteStream(delete); err != nil {
		return err
	}
	// poll until the stream is actually deleted
	describe := &kinesis.DescribeStreamInput{
		StreamName: aws.String(name),
	}
	// poll until the stream is actually created
	for i := 1; i < 1000; i++ {
		_, err := k.DescribeStream(describe)
		if err != nil {
			if awse, ok := err.(awserr.Error); ok && awse.Code() == "ResourceNotFoundException" {
				return nil
			}
			// Some other error
			return err
		}

		time.Sleep(1 * time.Millisecond)
	}
	return errors.New("timed out waiting for stream to disappear")
}

func TestClient() *kinesis.Kinesis {
	// Since we're communicating with a local kinesalite, we don't
	// need real credentials, but the AWS APIs expect us to have
	// nonempty values, so we have to fake them.
	creds := credentials.NewStaticCredentials("id", "secret", "token")
	cfg := aws.NewConfig().WithEndpoint(addr).WithRegion("local").WithCredentials(creds)
	svc := kinesis.New(session.New(cfg))

	// Add a handler which introduces a bit of delay to reading
	// responses. This is required because Kinesalite creates objects
	// like ShardIterators in callbacks _after_ it has responded to
	// requests. We need to wait for it to do all of its tasks.
	svc.Handlers.Unmarshal.PushFront(func(*request.Request) {
		time.Sleep(cbSleep)
	})

	return svc
}
