package events

import (
	"bytes"
	"context"
	"errors"
	"sync"
	"time"

	"code.justin.tv/sse/malachai/pkg/internal/closer"
	"code.justin.tv/sse/malachai/pkg/internal/fhevent"
	"code.justin.tv/sse/malachai/pkg/internal/fhevent/fheventiface"
	"code.justin.tv/sse/malachai/pkg/internal/stats/statsiface"
	"code.justin.tv/sse/malachai/pkg/log"
	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/awserr"
	"github.com/aws/aws-sdk-go/aws/session"
	"github.com/aws/aws-sdk-go/service/firehose"
)

// batchShipper builds batches for firehose then ships them
type batchShipper struct {
	Closer      *closer.Closer
	EventWriter fheventiface.EventWriterAPI
	Logger      log.S2SLogger

	// number of items we will buffer in memory
	MaxItemsInBuffer int

	// rate that batch will be flushed if not full
	// batch will also be flushed when hit either:
	// - item count exceeds firehoseMaxItemsInBatch
	// - size exceeds firehoseMaxBatchSizeInBytes
	BatchFlushRate time.Duration

	// rate that items will be flushed if not full
	// items are will also be flushed when they hit firehoseMaxRecordSizeInBytes
	ItemFlushRate time.Duration

	// firehose stream we will write to
	FirehoseDeliveryStreamName string

	// maximum time we will try to flush a single batch to firehose
	// if we hit this limit, items will be dropped.
	MaxFirehoseFlushTime time.Duration

	// exponential constant to use when retrying push if write fails
	FirehoseRetryRate time.Duration

	// TwitchTelemtry Client
	Statter statsiface.ReporterAPI

	maxItemsInBatch     int
	maxBatchSizeInBytes int

	runOnce         sync.Once
	initOnce        sync.Once
	initRan         bool
	batch           fhevent.Batch
	batchFlushTimer *time.Timer
	items           *fhevent.RecordBuilder
	itemFlushTimer  *time.Timer

	// triggers for unit tests
	onFlushIfNecessaryAndBatch func()
}

func (bs *batchShipper) Init(
	sess *session.Session,
	statter statsiface.ReporterAPI,
) (err error) {
	bs.initOnce.Do(func() {
		bs.Statter = statter
		bs.batchFlushTimer = time.NewTimer(bs.BatchFlushRate)
		bs.itemFlushTimer = time.NewTimer(bs.ItemFlushRate)
		bs.items = &fhevent.RecordBuilder{
			Item:               bytes.NewBuffer(nil),
			MaxItemSizeInBytes: fhevent.FirehoseMaxRecordSizeInBytes,
			ReadyChannel:       make(chan []byte, bs.MaxItemsInBuffer),
			Statter:            statter,
			Logger:             bs.Logger,
		}
		bs.EventWriter = &fhevent.Client{
			FirehoseAPI: firehose.New(sess),
			Logger:      bs.Logger,
		}
		bs.initRan = true
	})
	return
}

// Add an item to be built into a record.
func (bs *batchShipper) Add(ctx context.Context, item []byte) error {
	if !bs.initRan {
		return errors.New("Init() must be executed first")
	}
	return bs.items.Add(ctx, item)
}

// Run the batch shipper.
func (bs *batchShipper) Run() {
	bs.runOnce.Do(func() {
		// main loop
		bs.run(bs.Closer)

		// disable additions
		bs.items.Close()

		// flush out batch with timeout
		ctx, cancel := context.WithTimeout(context.Background(), bs.MaxFirehoseFlushTime)
		bs.finalFlush(ctx)
		cancel()
	})
}

func (bs *batchShipper) run(ctx context.Context) {
	for {
		if err := ctx.Err(); err != nil {
			return
		}
		bs.runLoop(ctx)
	}
}

func (bs *batchShipper) runLoop(ctx context.Context) {
	if (bs.batch.Len() >= bs.getMaxItemsInBatch()) || (bs.batch.SizeInBytes() >= bs.getMaxBatchSizeInBytes()) {
		ctx, cancel := context.WithTimeout(ctx, bs.MaxFirehoseFlushTime)
		bs.flushBatch(ctx)
		cancel()
		return
	}

	select {
	case <-ctx.Done():
		return
	case item := <-bs.items.ReadyChannel:
		ctx, cancel := context.WithTimeout(ctx, bs.MaxFirehoseFlushTime)
		bs.flushIfNecessaryAndBatch(ctx, item)
		cancel()
	case <-bs.itemFlushTimer.C:
		// items haven't been flushed in a while.
		ctx, cancel := context.WithTimeout(ctx, 100*time.Millisecond)
		bs.flushItem(ctx)
		cancel()
	case <-bs.batchFlushTimer.C:
		// batch hasn't been flushed in a while to firehose.
		ctx, cancel := context.WithTimeout(ctx, bs.MaxFirehoseFlushTime)
		bs.flushBatch(ctx)
		cancel()
	}
}

func (bs *batchShipper) finalFlush(ctx context.Context) {
	bs.Logger.Debugf("flushing last items with a timeout of: %s", bs.MaxFirehoseFlushTime)

	bs.flushBatch(ctx)

	if err := bs.items.Ship(ctx); err != nil {
		bs.Logger.Debugf("error in finalFlush while shipping logs to firehose: %s", err.Error())
		return
	}

	for {
		if done := bs.finalFlushLoop(ctx); done {
			return
		}
	}
}

func (bs *batchShipper) finalFlushLoop(ctx context.Context) (done bool) {
	if (bs.batch.Len() >= bs.getMaxItemsInBatch()) || (bs.batch.SizeInBytes() >= bs.getMaxBatchSizeInBytes()) {
		bs.flushBatch(ctx)
		return false
	}

	select {
	case <-ctx.Done():
		return true
	case item := <-bs.items.ReadyChannel:
		bs.flushIfNecessaryAndBatch(ctx, item)
	default:
		bs.flushBatch(ctx)
		return true
	}

	return false
}

func (bs *batchShipper) flushItem(ctx context.Context) {
	// regular flush of items from shipper failed
	// ignore error
	_ = bs.items.Ship(ctx)
	bs.itemFlushTimer.Reset(bs.ItemFlushRate)
}

func (bs *batchShipper) flushIfNecessaryAndBatch(ctx context.Context, item []byte) {
	expectedLen := bs.batch.Len() + 1
	expectedSize := bs.batch.SizeInBytes() + len(item)
	// if the batch is going to break limits, flushBatch first.
	if (expectedLen > bs.getMaxItemsInBatch()) || (expectedSize > bs.getMaxBatchSizeInBytes()) {
		if err := bs.flushBatchWithError(ctx); err != nil {
			bs.Logger.Debugf("error writing items to firehose: %s", err.Error())
			// optional signalling for unit tests
			if bs.onFlushIfNecessaryAndBatch != nil {
				bs.onFlushIfNecessaryAndBatch()
			}
			return
		}
	}

	bs.batch.Add(item)

	// optional signalling for unit tests
	if bs.onFlushIfNecessaryAndBatch != nil {
		bs.onFlushIfNecessaryAndBatch()
	}
}

func (bs *batchShipper) flushBatch(ctx context.Context) {
	if err := bs.flushBatchWithError(ctx); err != nil {
		bs.Logger.Debugf("error writing items to firehose: %s", err.Error())
	}
}

func (bs *batchShipper) flushBatchWithError(ctx context.Context) error {
	if bs.batch.Len() == 0 {
		return nil
	}

	if err := bs.EventWriter.PutRecordBatchRetriedWithContext(ctx, &firehose.PutRecordBatchInput{
		Records:            bs.batch.Batch(),
		DeliveryStreamName: aws.String(bs.FirehoseDeliveryStreamName),
	}); err != nil {
		if awsErr, ok := err.(awserr.Error); ok {
			if awsErr.Code() == firehose.ErrCodeInvalidArgumentException {
				// dont send invalid firehose record batches over and over again
				// Drop the batch and reset the flush timer.
				bs.batchFlushTimer.Reset(bs.BatchFlushRate)
				bs.batch.Clear()
				return nil
			}
		}
		bs.batchFlushTimer.Reset(bs.BatchFlushRate)
		return err
	}

	bs.batch.Clear()

	bs.batchFlushTimer.Reset(bs.BatchFlushRate)
	return nil
}

func (bs *batchShipper) getMaxItemsInBatch() int {
	if bs.maxItemsInBatch > 0 {
		return bs.maxItemsInBatch
	}
	return fhevent.FirehoseMaxItemsInBatch
}

func (bs *batchShipper) getMaxBatchSizeInBytes() int {
	if bs.maxBatchSizeInBytes > 0 {
		return bs.maxBatchSizeInBytes
	}
	return fhevent.FirehoseMaxBatchSizeInBytes
}
