package fhevent

import (
	"context"
	"errors"
	"time"

	"code.justin.tv/sse/malachai/pkg/backoff"
	"code.justin.tv/sse/malachai/pkg/log"
	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/service/firehose"
	"github.com/aws/aws-sdk-go/service/firehose/firehoseiface"
)

// firehose constants
const (
	// how big each record can be.
	FirehoseMaxRecordSizeInBytes = 500 * 1024 // 500KB
	// how many items we are able to send to firehose in each batch
	FirehoseMaxItemsInBatch = 500
	// how big each batch can be
	FirehoseMaxBatchSizeInBytes = 4 * 1024 * 1024 // 1MB
)

// Client is a covenience wrapper around firehose
type Client struct {
	firehoseiface.FirehoseAPI
	RetryRate time.Duration
	Logger    log.S2SLogger
}

var (
	errBackoff = errors.New("firehose backoff")
)

// PutRecordBatchRetriedWithContext puts a record batch all records are
// accounted for.
func (c *Client) PutRecordBatchRetriedWithContext(ctx context.Context, input *firehose.PutRecordBatchInput) error {
	er := backoff.ExponentialRetry{
		MaxTries: -1,
		Constant: c.RetryRate,
	}

	err := c.putRecordBatch(ctx, input)
	if err != errBackoff {
		return err
	}

	for er.Try(ctx, err) {
		err = c.putRecordBatch(ctx, input)
		if err != errBackoff {
			return err
		}
	}

	return er.Err()
}

func (c *Client) putRecordBatch(ctx context.Context, input *firehose.PutRecordBatchInput) error {
	res, err := c.FirehoseAPI.PutRecordBatchWithContext(ctx, input)
	// if we get a top level error, there is nothing we can do about this
	if err != nil {
		return err
	}

	if aws.Int64Value(res.FailedPutCount) == 0 {
		return nil
	}

	var unavailableErrors, invalidRecords, resourceNotFound int

	// go through each result and figure out what we need to resend
	records := input.Records[:0]
	for nResult, result := range res.RequestResponses {
		if result.ErrorCode != nil {
			switch aws.StringValue(result.ErrorCode) {
			case firehose.ErrCodeInvalidArgumentException:
				invalidRecords++
			case firehose.ErrCodeResourceNotFoundException:
				resourceNotFound++
			case firehose.ErrCodeServiceUnavailableException:
				unavailableErrors++
				fallthrough
			default:
				records = append(records, input.Records[nResult])
			}
		}

		if len(records) == int(aws.Int64Value(res.FailedPutCount)) {
			break
		}
	}

	if invalidRecords > 0 {
		c.Logger.Errorf("firehose reported %d invalid records", invalidRecords)
	}

	if resourceNotFound > 0 {
		c.Logger.Errorf("firehose reported %d resource not found records", resourceNotFound)
	}

	if len(records) == 0 {
		return nil
	}

	input.Records = records
	if unavailableErrors > 0 {
		return errBackoff
	}

	return nil
}
