package cloudwatch

import (
	"context"
	"sync"
	"time"

	"code.justin.tv/amzn/TwitchLogging"
	"code.justin.tv/amzn/TwitchProcessIdentifier"
	"code.justin.tv/amzn/TwitchTelemetry"
	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/session"
	"github.com/aws/aws-sdk-go/private/protocol/query"
	"github.com/aws/aws-sdk-go/service/cloudwatch"
	"github.com/aws/aws-sdk-go/service/cloudwatch/cloudwatchiface"
	"log"
	"math"
)

const (
	// CWMaxDatums represents the maximum datums per CW call
	// See: https://w.amazon.com/index.php/MonitoringTeam/CloudWatchFrontend/CloudWatchFrontEnd/PutMetricData
	CWMaxDatums = 1000
	// Do not retry
	maxRetries = 0
	// Timeout in seconds for requests to CloudWatch (after timeoutSeconds, the request will be cancelled)
	timeoutSeconds = 2
	// maxValuesPerDatum represents the maximum number of unique values we can have for any datum
	maxValuesPerDatum = 5000
)

const (
	// bufferedFlushPeriod represents the minimum time between auto-flushes by the Buffered CloudWatch sender
	// (used by Lambda). The Buffered sender will not auto-flush unless at least bufferedFlushPeriod has elapsed since
	// the last Flush() occurred.
	bufferedFlushPeriod = 30 * time.Second
	// bufferedTickerTime represents how often to check if an auto-flush is needed
	bufferedTickerTime = 5 * time.Second
)

// Buffered represents a buffered CloudWatch sender (commonly used with a Telemetry.SampleObserver)
type Buffered struct {
	*CommonSender
	aggregator  *telemetry.Aggregator
	stopSignal  chan bool
	stopWait    sync.WaitGroup
	flushPeriod time.Duration

	mux           sync.Mutex
	nextAutoFlush time.Time
}

// Unbuffered represents an unbuffered CloudWatch sender (commonly used with a Telemetry.BufferedAggregator)
type Unbuffered struct {
	*CommonSender
}

// CommonSender contains the common struct components and functions for MWS senders
type CommonSender struct {
	client            cloudwatchiface.CloudWatchAPI
	processIdentifier *identifier.ProcessIdentifier
	logger            logging.Logger
}

// New creates a new CloudWatch client configured using the given ProcessIdentifier
func New(processIdentifier *identifier.ProcessIdentifier, logger logging.Logger) telemetry.SampleObserver {
	sess := createAwsSession(processIdentifier.Region)
	return NewWithSession(processIdentifier, logger, sess)
}

func NewWithSession(processIdentifier *identifier.ProcessIdentifier, logger logging.Logger, sess *session.Session) telemetry.SampleObserver {
	cwClient := createCloudWatchClient(sess.Copy())

	common := &CommonSender{
		client:            cwClient,
		processIdentifier: processIdentifier,
		logger:            logger,
	}

	cw := &Buffered{
		CommonSender:  common,
		aggregator:    telemetry.NewAggregator(time.Minute),
		stopSignal:    make(chan bool),
		nextAutoFlush: time.Now(),
		flushPeriod:   bufferedFlushPeriod,
	}

	cw.run()
	return cw
}

// run flushes aggregated sample data to CloudWatch every bufferedTickerTime
func (cwBuff *Buffered) run() {
	cwBuff.stopWait.Add(1)
	go func() {
		ticker := time.NewTicker(bufferedTickerTime)
		defer ticker.Stop()
		for {
			select {
			case <-ticker.C:
				cwBuff.sharedFlush(true)
			case <-cwBuff.stopSignal:
				cwBuff.Flush()
				cwBuff.stopWait.Done()
				return
			}
		}
	}()
}

// Stop terminates the wait and flush loop that continuously sends metrics to CW on a set interval
// This is a synchronous operation, and current metrics will be flushed before this method returns
func (cwBuff *Buffered) Stop() {
	// send into the channel to initiate the stop
	cwBuff.stopSignal <- true
	// wait for the client to flush before returning
	cwBuff.stopWait.Wait()
}

// ObserveSample takes an internal collection of samples and sends it to CloudWatch
func (cwBuff *Buffered) ObserveSample(sample *telemetry.Sample) {
	// Ingest the sample into the Aggregator
	if cwBuff.aggregator != nil {
		cwBuff.aggregator.AggregateSample(sample)
	}
}

// Flush synchronously writes all aggregated samples to CloudWatch, wipes the
// aggregator clean, and returns
func (cwBuff *Buffered) Flush() {
	cwBuff.sharedFlush(false)
}

// sharedFlush is used to allow a single lock to handle flushing for both the direct (Flush) and indirect (poller) route
func (cwBuff *Buffered) sharedFlush(isBackground bool) {
	// Lock the world, since flush can be called manually or via run(). This ensures Lambda waits for these calls to
	// complete.
	cwBuff.mux.Lock()
	defer cwBuff.mux.Unlock()

	// Ensure we do not auto-flush if we last flushed less than flushPeriod ago
	// In a worst case scenario, this could mean auto-flushing once in flushPeriod + bufferedTickerTime
	if isBackground && time.Now().Before(cwBuff.nextAutoFlush) {
		return
	}

	// Update the next auto-flush timestamp so we won't auto-flush for at least another flushPeriod
	cwBuff.nextAutoFlush = time.Now().Add(cwBuff.flushPeriod)

	// Pull the distributions from the sample aggregator
	if cwBuff.aggregator != nil {
		distributions := cwBuff.aggregator.Flush()
		cwBuff.commonFlush(distributions)
	}
}

// NewUnbuffered returns a new CloudWatch client that needs to be run within a separate BufferedAggregator as it does not
// main a buffer/ any aggregation or flushing logic on its own
func NewUnbuffered(processIdentifier *identifier.ProcessIdentifier, logger logging.Logger) telemetry.SampleUnbufferedObserver {
	sess := createAwsSession(processIdentifier.Region)
	return NewUnbufferedWithSession(processIdentifier, logger, sess)
}

func NewUnbufferedWithSession(processIdentifier *identifier.ProcessIdentifier, logger logging.Logger, sess *session.Session) telemetry.SampleUnbufferedObserver {
	cwClient := createCloudWatchClient(sess.Copy())

	common := &CommonSender{
		client:            cwClient,
		processIdentifier: processIdentifier,
		logger:            logger,
	}

	return &Unbuffered{
		CommonSender: common,
	}
}

// FlushWithoutBuffering flushes without buffering the flush
func (cwUnbuff *Unbuffered) FlushWithoutBuffering(distributions []*telemetry.Distribution) {
	cwUnbuff.commonFlush(distributions)
}

// commonFlush is a helper function to runs the same functionality for both buffered and non-buffered flushes
func (cw *CommonSender) commonFlush(distributions []*telemetry.Distribution) {
	// Convert the resulting distributions into usable datums
	cwCollection := cw.constructMetricDatums(distributions)
	if len(cwCollection) == 0 {
		return
	}

	// We want to split metrics every CWMaxDatums elements.
	for i := 0; i < len(cwCollection); i += CWMaxDatums {
		// Determine the next endpoint
		colEnd := i + CWMaxDatums
		if colEnd > len(cwCollection) {
			colEnd = len(cwCollection)
		}

		// Send the sub-collection
		err := cw.sendMetrics(cwCollection[i:colEnd])
		if err != nil {
			// Check if the request was too large...
			if err == requestSizeError {
				// If so, split the request further by halving it until it's small enough or a single element remains
				cw.halveLargeRequest(cwCollection[i:colEnd])
			} else {
				cw.logFailedMetricsPublishIfNecessary(err)
			}
		}
	}
}

// halveLargeRequest allows us to break up requests that were found to be too large by halving the collection size and
// trying again, calling breakupLargeRequests to further break down the request until it is either small enough or a
// single datum remains
func (cw *CommonSender) halveLargeRequest(cwCollection []*cloudwatch.MetricDatum) {
	// If our length is 1, we don't split further. Technically, it is possible for this to be greater than the
	// request limit (since we are sending two arrays of float64 of max size 5000), but this is incredibly unlikely
	// since all values are likely to have been compressed by this point.
	if len(cwCollection) == 1 {
		cw.logFailedMetricsPublishIfNecessary(cw.sendMetrics(cwCollection))
		return
	}

	// Determine the step size (half the length of the list)
	halfStep := int(math.Ceil(float64(len(cwCollection)) / 2.0))
	for i := 0; i < len(cwCollection); i += halfStep {
		// Determine the next endpoint
		colEnd := i + halfStep
		if colEnd > len(cwCollection) {
			colEnd = len(cwCollection)
		}

		// Send the sub-collection
		err := cw.sendMetrics(cwCollection[i:colEnd])
		if err != nil {
			// Check if the request was too large...
			if err == requestSizeError {
				// If so, split the request further
				cw.halveLargeRequest(cwCollection[i:colEnd])
			} else {
				cw.logFailedMetricsPublishIfNecessary(err)
			}
		}
	}
}

// logFailedMetricsPublishIfNecessary logs a failed publish metrics request (if the request failed)
func (cw *CommonSender) logFailedMetricsPublishIfNecessary(err error) {
	if err == nil {
		return
	}
	if cw.logger != nil {
		cw.logger.Log("Failed to publish metrics", "err", err.Error())
	} else {
		log.Printf("Failed to publish metrics: %v", err)
	}
}

// sendMetrics performs the CloudWatch PutMetricData API call
func (cw *CommonSender) sendMetrics(cwCollection []*cloudwatch.MetricDatum) error {
	if len(cwCollection) == 0 {
		return nil
	}

	input := &cloudwatch.PutMetricDataInput{
		MetricData: cwCollection,
		// The namespace should just be the name of the service
		Namespace: aws.String(cw.processIdentifier.Service),
	}

	// Set a timeout for the CloudWatch call
	ctx, cancel := context.WithTimeout(context.Background(), timeoutSeconds*time.Second)
	defer cancel()
	_, err := cw.client.PutMetricDataWithContext(ctx, input)
	return err
}

// constructMetricDatums creates CloudWatch metric datums from the internal distributions
func (cw *CommonSender) constructMetricDatums(distributions []*telemetry.Distribution) []*cloudwatch.MetricDatum {
	allDatum := []*cloudwatch.MetricDatum{}
	for _, element := range distributions {
		newDatums := cw.constructNewDatums(element)
		if len(newDatums) == 0 {
			continue
		}
		allDatum = append(allDatum, newDatums...)
	}
	return allDatum
}

// constructNewDatums constructs a new CloudWatch metric datum from an internal distribution
func (cw *CommonSender) constructNewDatums(distribution *telemetry.Distribution) []*cloudwatch.MetricDatum {
	// Ensure there's a metric name
	if distribution.MetricID.Name == "" {
		return nil
	}
	// A single telemetry distribution has all related dimensions. These need to be broken into separate CW
	// datums to allow for rollups (since CW does not do this automatically).
	// This means n dimensions in the dimension list for one internal metric = n separate CW datums
	cwDatums := []*cloudwatch.MetricDatum{}

	// Create the single CW datum that will serve as the base for every other datum
	baseDatum := cloudwatch.MetricDatum{}

	// CW now supports values and counts. Use the histogram to construct these if possible
	// Each datum can have a maximum of maxValuesPerDatum unique values. We therefore need to make sure we break up the
	// values as they come along and construct separate datums if necessary
	var allValues [][]*float64
	var allCounts [][]*float64
	if distribution.SEH1 != nil {
		allValuesIdx := 0
		allValues = append(allValues, []*float64{})
		allCounts = append(allCounts, []*float64{})

		// Set the zero bucket if required
		if distribution.SEH1.ZeroBucket != 0 {
			allValues[allValuesIdx] = append(allValues[allValuesIdx], aws.Float64(0.0))
			allCounts[allValuesIdx] = append(allCounts[allValuesIdx], aws.Float64(float64(distribution.SEH1.ZeroBucket)))
		}
		// Set the individual values
		for bucketVal, sampleCount := range distribution.SEH1.Histogram {
			// As noted, there cannot be more than maxValuesPerDatum unique values per datum
			if len(allValues[allValuesIdx]) >= maxValuesPerDatum {
				allValuesIdx++
				allValues = append(allValues, []*float64{})
				allCounts = append(allCounts, []*float64{})
			}
			approximateVal := distribution.SEH1.ApproximateOriginalValue(bucketVal)
			allValues[allValuesIdx] = append(allValues[allValuesIdx], aws.Float64(approximateVal))
			allCounts[allValuesIdx] = append(allCounts[allValuesIdx], aws.Float64(float64(sampleCount)))
		}
	}
	// Always include a statistical set
	statisticalSet := &cloudwatch.StatisticSet{}
	statisticalSet.SetMaximum(distribution.Maximum)
	statisticalSet.SetMinimum(distribution.Minimum)
	statisticalSet.SetSampleCount(float64(distribution.SampleCount))
	statisticalSet.SetSum(distribution.Sum)
	baseDatum.SetStatisticValues(statisticalSet)

	// Set the common datum details
	baseDatum.SetMetricName(distribution.MetricID.Name)
	baseDatum.SetTimestamp(distribution.Timestamp.UTC())
	baseDatum.SetUnit(distribution.Unit)

	// If the length of allValues > 0, we had a usable SEH1, so we need to construct 1 or more defaultMetric depending
	// on the number of unique values our histogram had
	if len(allValues) > 0 {
		for allValIdx, valueList := range allValues {
			cwDatums = append(cwDatums, cw.createRollupDatums(baseDatum, distribution.MetricID.Dimensions, distribution.RollupDimensions, valueList, allCounts[allValIdx])...)
		}
	} else {
		// Use the baseDatum to construct all new datums, but don't set any values or counts
		cwDatums = append(cwDatums, cw.createRollupDatums(baseDatum, distribution.MetricID.Dimensions, distribution.RollupDimensions, nil, nil)...)
	}
	return cwDatums
}

// createRollupDatums constructs the original datum and all rollup datums
func (cw *CommonSender) createRollupDatums(baseDatum cloudwatch.MetricDatum, allDims map[string]string, distRollupDims [][]string, valueList []*float64, countsList []*float64) []*cloudwatch.MetricDatum {

	datumList := []*cloudwatch.MetricDatum{}

	// Start by determining the base dimensions (these will be used for a single datum)
	defaultDims := cw.constructDimensions(allDims)

	// Use the baseDatum to construct all new datums
	defaultMetric := baseDatum
	defaultMetric.SetDimensions(defaultDims)
	if len(valueList) > 0 {
		defaultMetric.SetValues(valueList)
	}
	if len(countsList) > 0 {
		defaultMetric.SetCounts(countsList)
	}
	datumList = append(datumList, &defaultMetric)

	// Construct the appropriate roll ups
	for _, rollupDims := range distRollupDims {
		rollupDim := cw.createRollupDimensions(allDims, rollupDims)
		rollupMetric := baseDatum
		rollupMetric.SetDimensions(rollupDim)
		if len(valueList) > 0 {
			rollupMetric.SetValues(valueList)
		}
		if len(countsList) > 0 {
			rollupMetric.SetCounts(countsList)
		}
		datumList = append(datumList, &rollupMetric)
	}
	return datumList
}

// createRollupDimensions takes a map of dimensions and rollupDimensions and constructs the appropriate CW dimensions
func (cw *CommonSender) createRollupDimensions(allDims map[string]string, rollups []string) []*cloudwatch.Dimension {
	// Copy the original map
	dims := make(telemetry.DimensionSet)
	for name, value := range allDims {
		dims[name] = value
	}

	// Remove all the rollup dimensions from the original dimension map
	for _, name := range rollups {
		delete(dims, name)
	}
	// Convert the updated dimension map into CW dimensions
	return cw.constructDimensions(dims)
}

// constructDimensions takes a map of dimensions (string->string) and converts it into a list of CW dimensions
func (cw *CommonSender) constructDimensions(dimensions telemetry.DimensionSet) []*cloudwatch.Dimension {
	cwDims := []*cloudwatch.Dimension{}
	for name, value := range dimensions {
		// Ensure the dimensions are valid (require both a name and a value)
		if (name == "") || (value == "") {
			continue
		}

		cwDim := cloudwatch.Dimension{}
		cwDim.SetName(name)
		cwDim.SetValue(value)
		cwDims = append(cwDims, &cwDim)
	}

	return cwDims
}

func createAwsSession(region string) *session.Session {
	return session.Must(session.NewSessionWithOptions(session.Options{
		Config: aws.Config{
			Region:     aws.String(region),
			MaxRetries: aws.Int(maxRetries),
		},
	}))
}

func createCloudWatchClient(sess *session.Session) *cloudwatch.CloudWatch {
	// Construct a default CW client
	cwClient := cloudwatch.New(sess)
	// Remove the default BuildHandler
	cwClient.Handlers.Build.Remove(query.BuildHandler)
	// Add the GZip handler
	cwClient.Handlers.Build.PushBack(buildPostGZip)
	return cwClient
}
