package dynamo

import (
	"strconv"

	"code.justin.tv/creator-collab/log/errors"
	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/service/dynamodb"
	"github.com/aws/aws-sdk-go/service/dynamodb/dynamodbattribute"
	"golang.org/x/net/context"
)

// If we ever write to more than 25 buckets we'll need to update these methods to
// make several batch calls since we can only get or write 25 items at a time.
//
// ChannelIDs in this case are user logins, not their numeric IDs. The maximum
// length on a login is 25 characters. Assumming logins of 15 characters we can store
// 26,667 channelIDs per bucket.
//
// As of 6/14/2018 we have 45,000 - 50,000 live channels at any time. This
// current value of 8 should accommodate more than 200,000 live channels.
const liveUsersNumberOfBuckets = 8

// CacheLiveUsers saves a snapshot of live channel IDs. DynamoDB limits each item to
// 400KB, which is not enough to store all the live channel IDs. We split the IDs
// across several buckets so that we can store them all.
func (db *Dynamo) CacheLiveUsers(ctx context.Context, channelIDs map[string]bool) error {
	// Create a 2 dimensional slice of strings to store all the channel IDs
	buckets := make([][]string, liveUsersNumberOfBuckets)
	for i := range buckets {
		buckets[i] = []string{}
	}

	for channelID := range channelIDs {
		bucketNumber := GetBucketNumber(liveUsersNumberOfBuckets, channelID)
		buckets[bucketNumber] = append(buckets[bucketNumber], channelID)
	}

	// Create an input struct with an item for each bucket we want to write
	channelWriteRequests := []*dynamodb.WriteRequest{}
	for i := range buckets {
		// DynamoDB does not allow empty sets so we have to set the value to NULL
		var idsValue dynamodb.AttributeValue
		if len(buckets[i]) == 0 {
			idsValue = dynamodb.AttributeValue{NULL: aws.Bool(true)}
		} else {
			idsValue = dynamodb.AttributeValue{SS: aws.StringSlice(buckets[i])}
		}

		channelWriteRequests = append(channelWriteRequests, &dynamodb.WriteRequest{
			PutRequest: &dynamodb.PutRequest{
				Item: map[string]*dynamodb.AttributeValue{
					"IDs": &idsValue,
					"Key": {
						S: aws.String("LiveChannels"),
					},
					"Bucket": {
						N: aws.String(strconv.Itoa(i)),
					},
				},
			},
		})
	}

	batchInput := &dynamodb.BatchWriteItemInput{
		RequestItems: map[string][]*dynamodb.WriteRequest{
			"Caches": channelWriteRequests,
		},
	}

	_, err := db.client.BatchWriteItemWithContext(ctx, batchInput)
	if err != nil {
		return errors.Wrap(err, "dynamodb - caching live users failed")
	}

	return nil
}

// GetCachedLiveUsers gets the list of live users from the current cache. It
// reads from all the buckets and returns a single map to approximate a set data
// type. This method uses BatchGetItem rather than a Query because Queries can
// return at most 1MB of data with a single operation. BatchGetItem can return
// up to 16MB of data across 100 items in a single operation.
func (db *Dynamo) GetCachedLiveUsers(ctx context.Context) (map[string]bool, error) {
	// Create an item map for each bucket that we want to query
	itemKeys := []map[string]*dynamodb.AttributeValue{}
	for i := 0; i < liveUsersNumberOfBuckets; i++ {
		itemKeys = append(itemKeys, map[string]*dynamodb.AttributeValue{
			"Key": {
				S: aws.String("LiveChannels"),
			},
			"Bucket": {
				N: aws.String(strconv.Itoa(i)),
			},
		})
	}

	input := &dynamodb.BatchGetItemInput{
		RequestItems: map[string]*dynamodb.KeysAndAttributes{
			"Caches": {
				Keys:                 itemKeys,
				ProjectionExpression: aws.String("IDs"),
			},
		},
	}

	result, err := db.client.BatchGetItemWithContext(ctx, input)
	if err != nil {
		return nil, errors.Wrap(err, "dynamodb - getting cached live users failed")
	}

	// Build the map by reading each ID set from the results
	channelIDs := make(map[string]bool)
	for _, cache := range result.Responses["Caches"] {
		ids := []string{}
		err = dynamodbattribute.Unmarshal(cache["IDs"], &ids)
		if err != nil {
			return nil, errors.Wrap(err, "dynamodb - getting cached live users - unmarshal failed")
		}
		for _, id := range ids {
			channelIDs[id] = true
		}
	}

	return channelIDs, nil
}
