package sink

import (
	"fmt"
	"log"
	"os"
	"sync"
	"time"

	"code.justin.tv/ear/kinesis/combiner"
	"code.justin.tv/release/trace/internal/awsutils"
	"code.justin.tv/release/trace/internal/pipeline"
	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/session"
	"github.com/aws/aws-sdk-go/service/kinesis"
	"github.com/cactus/go-statsd-client/statsd"
	"golang.org/x/net/context"
)

const (
	kinesisBatchMaxRecords = 500

	defaultFlushInterval = time.Millisecond * 100
)

type kinesisSink struct {
	client *kinesis.Kinesis
	cfg    KinesisConfig
	buffer []*combiner.Blob
	size   int
	wg     sync.WaitGroup
}

type KinesisConfig struct {
	// Name of stream to connect to.
	Stream string

	// AWS Region to connect to. Defaults to us-west-2.
	Region string

	// Callback that will be called in case of errors. Defaults to
	// ignoring all errors.
	OnError pipeline.ErrorHandler

	// Receiver for log messages. Defaults to logging to stderr.
	Logger *log.Logger

	// Receiver for statistics. Defaults to a Noop statter. The following
	// counters are written if a real statter is provided:
	//
	//    putrecords.attempted
	//      Number of API calls to write a batch of records to Kinesis
	//
	//    putrecords.errors
	//      Number of calls to the PutRecords API which were rejected
	//      completely for having errors.
	//
	//    records_flushed
	//      Number of records flushed out of internal buffer and to
	//      Kinesis. Retried records are counted just once.
	//
	//    records_sent
	//      Number of records sent to Kinesis. Retried records are counted again
	//      for each time they are resent.
	//
	//    records_succeeded, byshard.<shardid>.records_succeeded
	//      Number of records accepted by Kinesis.
	//
	//    records_failed.throttled, byshard.<shardid>.records_failed.throttled
	//      Number of records which were rejected by Kinesis for exceeding
	//      provisioned throughput.
	//
	//    records_failed.internal_error, byshard.<shardid>.records_failed.internal_error
	//      Number of records rejected by Kinesis because of a failure internal
	//      to Kinesis.
	//
	//    records_failed.unknown_reason, byshard.<shardid>.records_failed.unknown_reason
	//      Number of records rejected by Kinesis without a known explanation.
	//
	//    records_dropped
	//      Number of records rejected by Kinesis enough times that they
	//      had to be dropped and never got written.
	//
	//
	//  Also, the following timers are written:
	//
	//    putrecords
	//      Time in milliseconds spent in PutRecords calls.
	Stats statsd.Statter

	// Maximum number of retries when attempting to put a batch into
	// kinesis. Defaults to 5
	MaxRetries int

	// The exponential backoff unit for how long to wait between retries.
	RetryDelay time.Duration

	// The maximum amount of time to wait between retries
	MaxBackoff time.Duration

	// Maximum size of a batch in bytes. Defaults to 4.9 MB.
	BatchSizeBytes int

	// How often to flush to kinesis. Defaults to 100ms.
	FlushInterval time.Duration
}

func NewKinesisConfig(stream string) KinesisConfig {
	stats, _ := statsd.NewNoopClient()
	return KinesisConfig{
		Stream:         stream,
		Region:         "us-west-2",
		OnError:        pipeline.IgnoreErrors,
		Logger:         log.New(os.Stderr, "[kinesis] ", log.LstdFlags|log.Lmicroseconds|log.Lshortfile),
		MaxRetries:     5,
		RetryDelay:     100 * time.Millisecond,
		MaxBackoff:     5 * time.Second,
		BatchSizeBytes: 4900000, // 4.9 MB
		FlushInterval:  defaultFlushInterval,
		Stats:          stats,
	}
}

func Kinesis(ctx context.Context, cfg KinesisConfig) (Sink, error) {
	awsCfg := aws.NewConfig().WithRegion(cfg.Region)
	awsCfg = awsutils.WithChitin(ctx, awsCfg)

	client := kinesis.New(session.New(), awsCfg)
	sink := &kinesisSink{
		cfg:    cfg,
		client: client,
		buffer: []*combiner.Blob{},
	}
	return sink.Consume, nil
}

func (ks *kinesisSink) Consume(data <-chan *combiner.Blob) {
	flusher := time.NewTicker(ks.cfg.FlushInterval)

RunLoop:
	for {
		select {
		case <-flusher.C:
			ks.flush()
		case b, open := <-data:
			if !open {
				ks.flush()
				break RunLoop
			}
			s := len(b.Data)
			if s+ks.size > ks.cfg.BatchSizeBytes || len(ks.buffer) == kinesisBatchMaxRecords {
				ks.flush()
			}
			ks.buffer = append(ks.buffer, b)
			ks.size += s
		}
	}

	flusher.Stop()
	ks.wg.Wait()
}

func (ks *kinesisSink) flush() {
	// Skip if there's nothing to flush
	if len(ks.buffer) == 0 {
		return
	}

	// Build the args while we can, because blobs will shortly change from underneath us
	recs := make([]*kinesis.PutRecordsRequestEntry, len(ks.buffer))
	for i, b := range ks.buffer {
		recs[i] = &kinesis.PutRecordsRequestEntry{
			PartitionKey:    aws.String(b.Key),
			ExplicitHashKey: aws.String(b.Key),
			Data:            b.Data,
		}
	}
	ks.buffer = ks.buffer[:0]
	ks.size = 0

	ks.wg.Add(1)
	go ks.putRecords(recs)
}

func (ks *kinesisSink) putRecords(recs []*kinesis.PutRecordsRequestEntry) {
	defer ks.wg.Done()

	args := &kinesis.PutRecordsInput{
		StreamName: aws.String(ks.cfg.Stream),
		Records:    recs,
	}
	ks.cfg.Stats.Inc("records_flushed", int64(len(recs)), 1)

	for attempt := 1; attempt <= ks.cfg.MaxRetries; attempt++ {
		sleepFor := backoff(attempt-1, ks.cfg.RetryDelay, ks.cfg.MaxBackoff)
		time.Sleep(sleepFor)

		ks.cfg.Stats.Inc("putrecords.attempted", 1, 1)
		ks.cfg.Stats.Inc("records_sent", int64(len(recs)), 1)

		t0 := time.Now()
		res, err := ks.client.PutRecords(args)
		ks.cfg.Stats.TimingDuration("putrecords", time.Now().Sub(t0), 1)

		if err != nil {
			ks.cfg.Logger.Printf("PutRecords failure attempt number %d / %d: %v", attempt, ks.cfg.MaxRetries, err)
			ks.cfg.Stats.Inc("putrecords.errors", 1, 1)
			continue
		}

		// Find all failed records and update the slice to contain only failures
		i := 0
		for j, result := range res.Records {
			shard := aws.StringValue(result.ShardId)
			if shard == "" {
				shard = "unknown"
			}

			if aws.StringValue(result.ErrorCode) != "" {
				ks.cfg.Logger.Printf("Re-enqueueing failed record to buffer for retry (attempt #%d). error code=%q  message=%q", attempt, aws.StringValue(result.ErrorCode), aws.StringValue(result.ErrorMessage))
				switch aws.StringValue(result.ErrorCode) {
				case "ProvisionedThroughputExceededException":
					ks.cfg.Stats.Inc("records_failed.throttled", 1, 1)
					ks.cfg.Stats.Inc(fmt.Sprintf("byshard.%s.records_failed.throttled", shard), 1, 1)
				case "InternalFailure":
					ks.cfg.Stats.Inc("records_failed.internal_error", 1, 1)
					ks.cfg.Stats.Inc(fmt.Sprintf("byshard.%s.records_failed.internal_error", shard), 1, 1)
				default:
					// Something undocumented
					ks.cfg.Stats.Inc("records_failed.unknown_reason", 1, 1)
					ks.cfg.Stats.Inc(fmt.Sprintf("byshard.%s.records_failed.unknown_reason", shard), 1, 1)
				}
				args.Records[i] = args.Records[j]
				i++
			} else {
				ks.cfg.Stats.Inc("records_succeeded", 1, 1)
				ks.cfg.Stats.Inc(fmt.Sprintf("byshard.%s.records_succeeded", shard), 1, 1)
			}
		}
		args.Records = args.Records[:i]

		if len(args.Records) == 0 {
			// No records need to be retried.
			return
		}
	}
	// If we get here, we never returned from inside the loop, so we've failed :(
	err := fmt.Errorf("Failed sending %d out of %d records to kinesis", len(args.Records), len(recs))
	ks.cfg.Stats.Inc("records_dropped", int64(len(args.Records)), 1)
	ks.cfg.OnError(err)
}
