package stat

import (
	"sync"
	"time"

	"code.justin.tv/systems/sandstorm/logging"
	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/service/cloudwatch"
)

// thread safe batcher that stores batches of cloudwatch metrics
type batcher struct {
	Logger logging.Logger

	initSync sync.Once

	// counter metrics added by Increment
	stats       map[metricKey]*statisticSet
	statsAccess sync.Mutex
}

func (b *batcher) init() {
	b.initSync.Do(func() {
		b.popStats()
	})
}

func (b *batcher) currentTimeBucket(resolution time.Duration) timeBucket {
	return timeBucket(time.Now().UTC().Truncate(resolution))
}

func (b *batcher) popStats() map[metricKey]*statisticSet {
	b.statsAccess.Lock()
	stats := b.stats
	b.stats = make(map[metricKey]*statisticSet)
	b.statsAccess.Unlock()
	return stats
}

func (b *batcher) Increment(metric Metric) {
	b.init()

	b.withStatisticSetFor(metric, func(s *statisticSet) {
		s.Increment()
	})
}

func (b *batcher) Measure(metric Metric, value float64) {
	b.init()

	b.withStatisticSetFor(metric, func(s *statisticSet) {
		s.Record(value)
	})
}

func (b *batcher) withStatisticSetFor(metric Metric, fn func(*statisticSet)) {
	if metric.StorageResolution == time.Duration(0) {
		metric.StorageResolution = time.Minute
	}

	if time.Time(metric.timeBucket).IsZero() {
		metric.timeBucket = b.currentTimeBucket(metric.StorageResolution)
	}

	key, err := metric.toKey()
	if err != nil {
		b.Logger.Error(err)
		return
	}

	b.statsAccess.Lock()
	if _, ok := b.stats[key]; !ok {
		b.stats[key] = &statisticSet{}
	}
	fn(b.stats[key])
	b.statsAccess.Unlock()
}

func (b *batcher) Pop() []*cloudwatch.MetricDatum {
	b.init()

	stats := b.popStats()

	base := func() *cloudwatch.MetricDatum {
		return &cloudwatch.MetricDatum{
			Unit: aws.String(cloudwatch.StandardUnitCount),
		}
	}

	values := make([]*cloudwatch.MetricDatum, 0, len(stats))
	for metricKey, stat := range stats {
		var metric Metric
		if err := metric.fromKey(metricKey); err != nil {
			b.Logger.Errorf("error popping metricKey: %s", err)
		}

		values = append(values, base().
			SetStorageResolution(int64(metric.StorageResolution.Seconds())).
			SetDimensions(metric.Dimensions).
			SetMetricName(metric.MetricName).
			SetStatisticValues(stat.StatisticSet()).
			SetTimestamp(time.Time(metric.timeBucket)))
	}

	return values
}
