package telemetry

import (
	"code.justin.tv/amzn/TwitchLogging"
	"log"
	"sync"
	"sync/atomic"
	"time"
)

const (
	droppedMetricErrorMessage = "TwitchTelemetry BufferedAggregator forced to drop metrics. More info: https://docs.fulton.twitch.a2z.com/docs/metrics.html#dropped-metrics"
	droppedMetricKey          = "DroppedMetrics"
	timestampKey              = "Timestamp"
)

// BufferedAggregator keeps a distribution for each unique metric ingested
type BufferedAggregator struct {
	aggregationPeriod        time.Duration
	sampleChannelBufferSize  int
	sampleChannel            chan *Sample
	workToBeDoneStopSignal   chan bool
	flushPeriod              time.Duration
	sender                   SampleUnbufferedObserver
	stopWait                 sync.WaitGroup
	logger                   logging.Logger
	droppedMetrics           int32
	droppedMetricsStopSignal chan bool
}

// NewBufferedAggregator returns a new, empty buffered Aggregator. Unlike telemetry.Aggregator, this is buffered and
// dictates when the senders should flush metrics as opposed to the senders maintaining this logic.
//
// Set flushPeriod to the frequency you wish to flush metrics
// Set aggregationPeriod to the time resolution needed, for example time.Minute for one-minute metrics.
// Set aggregationPeriod to zero to disable this time bucketing and aggregate samples regardless of timestamp.
func NewBufferedAggregator(flushPeriod time.Duration, bufferSize int, aggregationPeriod time.Duration, sender SampleUnbufferedObserver, logger logging.Logger) SampleObserver {
	aggregator := &BufferedAggregator{
		aggregationPeriod:        aggregationPeriod,
		sampleChannelBufferSize:  bufferSize,
		sampleChannel:            make(chan *Sample, bufferSize),
		workToBeDoneStopSignal:   make(chan bool),
		droppedMetricsStopSignal: make(chan bool),
		droppedMetrics:           0,
		flushPeriod:              flushPeriod,
		sender:                   sender,
		logger:                   logger,
	}

	aggregator.runDroppedMetricsLogger()
	aggregator.runWorkToBeDone()
	return aggregator
}

// ObserveSample takes a sample and observes it by aggregating it into
// the distribution for the sample's metric. If the distribution
// does not exist, it will be created.
func (a *BufferedAggregator) ObserveSample(sample *Sample) {
	select {
	// Send the sample to a buffered channel
	case a.sampleChannel <- sample:
	default:
		// Note the number of failed metrics
		atomic.AddInt32(&a.droppedMetrics, 1)
	}
}

// Flush is not a supported operation, since we are using a buffer to handle all flushing
func (a *BufferedAggregator) Flush() {
	a.log("Flushing a buffered aggregator is not supported")
}

// Close stops the process buffer and ends metrics processing
func (a *BufferedAggregator) Stop() {
	a.workToBeDoneStopSignal <- true
	a.stopWait.Wait()
}

// runWorkToBeDone runs through the buffer
func (a *BufferedAggregator) runWorkToBeDone() {
	// If we lack required components, return
	if a.sender == nil || a.sampleChannelBufferSize <= 0 {
		a.log("Buffered aggregator must have a sender and positive buffer size")
		return
	}

	// Add a stopWait, this is solely used to ensure we block on the final close buffer channel (when we shut down
	// the service, we block to return our final metrics)
	a.stopWait.Add(1)

	go func() {
		flushTicker := time.NewTicker(a.flushPeriod)
		defer flushTicker.Stop()
		defer a.stopWait.Done()

		// distributions accumulates incoming samples between flushes.
		distributions := make(map[string]*Distribution)

		for {
			select {
			case sample := <-a.sampleChannel:
				// Process new samples as they come in
				a.processBuffer(distributions, sample)
			case <-flushTicker.C:
				// Flush when the ticker fires

				// Swap the current map out for a new one.
				flushedDistributions := distributions
				distributions = make(map[string]*Distribution)

				// Ownership of the old map pointed to by flushedDistributions is passed to a new goroutine here.
				// After this point, the old map isn't accessed by this one.
				// This goroutine then continues on to process more samples into the new map while the other
				// handles sending the old map that was given to it.
				go a.processDistribution(flushedDistributions)
			case <-a.workToBeDoneStopSignal:
				// Flush up to the entire buffer size (since there's no guarantee there aren't more buffer elements
				// that need to be processed)
				var done bool
				for i := 0; i < a.sampleChannelBufferSize && !done; i++ {
					select {
					case sample := <-a.sampleChannel:
						a.processBuffer(distributions, sample)
					default:
						done = true
					}
				}

				// Final flush to ensure all metrics have been emitted
				// Unlike above, this is done synchronously without a new goroutine so that it can block server shutdown.
				a.processDistribution(distributions)

				// Stop the dropped metrics logger
				// This happens after the sample channel is closed to ensure we don't have any additional metrics
				a.droppedMetricsStopSignal <- true
				return
			}
		}
	}()
}

// runDroppedMetricsLogger periodically logs whether any metrics were dropped
func (a *BufferedAggregator) runDroppedMetricsLogger() {
	a.stopWait.Add(1)
	go func() {
		ticker := time.NewTicker(30 * time.Second)
		defer ticker.Stop()
		defer a.stopWait.Done()
		for {
			select {
			case <-ticker.C:
				// Log dropped metrics if necessary
				a.logDroppedMetrics()
			case <-a.droppedMetricsStopSignal:
				// Log dropped metrics if necessary
				a.logDroppedMetrics()
				return
			}
		}
	}()
}

// logDroppedMetrics is a helper function to log failed metrics if necessary
func (a *BufferedAggregator) logDroppedMetrics() {
	// Get the number of failed metrics and reset the counter to 0
	droppedMetrics := atomic.SwapInt32(&a.droppedMetrics, 0)
	if droppedMetrics > 0 {
		if a.logger != nil {
			a.logger.Log(droppedMetricErrorMessage,
				droppedMetricKey, droppedMetrics,
				timestampKey, time.Now().UTC().Format(time.RFC3339))
		} else {
			log.Printf("%s %s %v. %s %s",
				droppedMetricErrorMessage,
				droppedMetricKey,
				droppedMetrics,
				timestampKey,
				time.Now().UTC().Format(time.RFC3339))
		}
	}
}

// processBuffer serves as a helper function to process samples in the buffered channel.
func (a *BufferedAggregator) processBuffer(distributionMap map[string]*Distribution, sample *Sample) {
	aggregationTimestamp := alignToAggregationPeriod(a.aggregationPeriod, sample.Timestamp)
	hashKey := uniqueKeyForMap(sample, aggregationTimestamp, a.aggregationPeriod)

	distribution, found := distributionMap[hashKey]
	if !found {
		distribution = NewDistributionFromSample(sample)
		distribution.Timestamp = aggregationTimestamp
		distributionMap[hashKey] = distribution
	} else {
		distribution.AddSample(sample)
	}
}

// processDistribution serves as a helper function to process the distribution
func (a *BufferedAggregator) processDistribution(distributionMap map[string]*Distribution) {
	distributions := make([]*Distribution, 0, len(distributionMap))
	for _, d := range distributionMap {
		distributions = append(distributions, d)
	}
	a.sender.FlushWithoutBuffering(distributions)
}

// log messages
func (a *BufferedAggregator) log(msg string) {
	if a.logger != nil {
		a.logger.Log(msg)
	} else {
		log.Printf(msg)
	}
}
