package telemetry

import (
	"sort"
	"strconv"
	"sync"
	"time"
)

// keySeparator is used to separate values when converting a metric id into a unique key for use in maps.
// Metric names and dimension names should not contain newlines.
var keySeparator = "\n"

// Aggregator keeps a distribution for each unique metric ingested
type Aggregator struct {
	aggregationPeriod time.Duration
	mux               sync.Mutex
	distributions     map[string]*Distribution
}

// NewAggregator returns a new, empty Aggregator.
// 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 NewAggregator(aggregationPeriod time.Duration) *Aggregator {
	return &Aggregator{
		distributions:     make(map[string]*Distribution),
		aggregationPeriod: aggregationPeriod,
	}
}

// AggregateSample takes a sample and aggregates it into
// the distribution for the sample's metric. If the distribution
// does not exist, it will be created.
func (a *Aggregator) AggregateSample(sample *Sample) {
	// TODO: this is potentially inefficient, but it's the easiest code that works (lock the world)
	a.mux.Lock()
	defer a.mux.Unlock()
	// TODO: Do we need a global struct lock? I feel like we do, but
	// it feels redundant to have a global struct lock to add samples while
	// _also_ have one lock per distribution...
	aggregationTimestamp := alignToAggregationPeriod(a.aggregationPeriod, sample.Timestamp)
	hashKey := uniqueKeyForMap(sample, aggregationTimestamp, a.aggregationPeriod)

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

// getDistributions returns all distributions in the Aggregator.
// Does not modify or clear the distributions
// Note that if the aggregator is not reset() after a call to getDistributions()
// then the distributions can continue to be modified by called to AggregateSample()
// Making this a private method avoids some unexpected concurrency issues
func (a *Aggregator) getDistributions() []*Distribution {
	distributions := make([]*Distribution, 0, len(a.distributions))
	for _, d := range a.distributions {
		distributions = append(distributions, d)
	}
	return distributions
}

// Reset wipes the Aggregator clean by deleting all Distributions
func (a *Aggregator) Reset() {
	a.mux.Lock()
	defer a.mux.Unlock()
	a.reset()
}

// assumes lock is acquired!
func (a *Aggregator) reset() {
	a.distributions = make(map[string]*Distribution)
}

// Flush returns current distributions, while also resetting them internally
func (a *Aggregator) Flush() []*Distribution {
	a.mux.Lock()
	defer a.mux.Unlock()

	distributions := a.getDistributions()
	a.reset()
	return distributions
}

func alignToAggregationPeriod(aggregationPeriod time.Duration, timestamp time.Time) time.Time {
	if aggregationPeriod > 0 {
		return timestamp.Truncate(aggregationPeriod)
	}
	return timestamp
}

func uniqueKeyForMap(sample *Sample, aggregationTimestamp time.Time, aggregationPeriod time.Duration) string {
	// Metrics with dimensions for Service, Stage, Substage, Region, and OperationName seem to be around 130 bytes.
	b := make([]byte, 0, 200)
	b = append(b, []byte(sample.MetricID.Name)...)
	if aggregationPeriod > 0 {
		timeKey := strconv.FormatInt(aggregationTimestamp.Unix(), 10)
		b = append(b, []byte(keySeparator)...)
		b = append(b, []byte(timeKey)...)
	}
	keys := make([]string, 0, len(sample.MetricID.Dimensions))
	for k := range sample.MetricID.Dimensions {
		keys = append(keys, k)
	}
	sort.Strings(keys)
	for _, k := range keys {
		b = append(b, []byte(keySeparator)...)
		b = append(b, []byte(k)...)
		b = append(b, []byte(keySeparator)...)
		b = append(b, []byte(sample.MetricID.Dimensions[k])...)
	}
	return string(b)
}
