/*
Package ddbmetrics sends DynamoDB provisioned and consumed read/write capacity units to statsd.

Every time you do a DynamoDB api call, you can request that the response includes how much capacity you consumed during the request.  ddbmetrics handles collecting those metrics and sending them to statsd for you.

Since a table's provisioned RCU/WCU values can be changed at any time, those values are collected from DynamoDB every 10 minutes. This allows you to generate percentage values in grafana.
*/
package ddbmetrics

import (
	"fmt"
	"log"
	"strconv"
	"sync"
	"time"

	"github.com/afex/hystrix-go/hystrix/rolling"
	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/service/dynamodb"
	"github.com/cactus/go-statsd-client/statsd"
)

const sampleRate = 1.0

// Operation represents the whether the capacity data reflects a read or write
type Operation int

const (
	// Read capacity
	Read Operation = iota
	// Write capacity
	Write
)

// Publisher keeps state of the last 10 seconds of consumed capacity for tables
type Publisher struct {
	db    *dynamodb.DynamoDB
	c     chan capacity
	stats statsd.Statter
	// TableDescriptions stores cached metadata about all tables consumption has been reported on.  Refreshed every 10 minutes.
	TableDescriptions map[string]*dynamodb.TableDescription
	usedCapacityMu    sync.Mutex
	usedCapacity      map[string]*rolling.Number
	tables            map[string]struct{}
	wg                sync.WaitGroup
	stop              chan struct{}
}

// capacity wraps ConsumedCapacity, including whether the operation was a read or write.
type capacity struct {
	consumed *dynamodb.ConsumedCapacity
	op       Operation
}

// New creates an idle Publisher, ready to receive capacity metrics.
func New(db *dynamodb.DynamoDB, stats statsd.Statter) *Publisher {
	return &Publisher{
		db:                db,
		stats:             stats,
		c:                 make(chan capacity),
		TableDescriptions: make(map[string]*dynamodb.TableDescription),
		tables:            make(map[string]struct{}),
		usedCapacity:      make(map[string]*rolling.Number),
		stop:              make(chan struct{}),
	}
}

// Start launches backround workers to process incoming metrics, and report provision/consumed capacity values
func (p *Publisher) Start() {
	p.pollProvisionedThroughput()
	p.measureUsedCapacity()
	p.recordUsedCapacity()
}

// Close signals and waits for running workers to stop
func (p *Publisher) Close() {
	close(p.stop)
	p.wg.Wait()
}

// Report add the given capacity into the 10 second rolling window for the relevant table.
func (p *Publisher) Report(rw Operation, cc *dynamodb.ConsumedCapacity) {
	p.c <- capacity{
		consumed: cc,
		op:       rw,
	}
}

// ReportBatch reports a list of capacities, used during batch operations.
func (p *Publisher) ReportBatch(rw Operation, cc []*dynamodb.ConsumedCapacity) {
	for _, c := range cc {
		p.Report(rw, c)
	}
}

// GetUsedCapacity returns the average for the last 10 seconds
func (p *Publisher) GetUsedCapacity(rw Operation, table string, indexName *string) float64 {
	n := p.capacityNumber(rw, table, indexName)
	return n.Avg(time.Now())
}

// pollProvisionedThroughput reads RCU/WCU settings for each table and index at an interval.
func (p *Publisher) pollProvisionedThroughput() {
	p.wg.Add(1)
	go func() {
		defer p.wg.Done()
		c := time.Tick(10 * time.Minute)
		for {
			select {
			case <-c:
				p.readProvisionedThroughput()
			case <-p.stop:
				return
			}
		}
	}()
}

// readProvisionedThroughput fetches table/index metadata and sends provisioned capacity metrics to graphite
func (p *Publisher) readProvisionedThroughput() {
	if p.stats != nil {
		for t := range p.tables {
			tableDescription := p.fetchTableDescription(t)
			if tableDescription != nil {
				p.TableDescriptions[t] = tableDescription
			}
		}

		for k, v := range p.TableDescriptions {
			p.stats.Gauge(fmt.Sprintf("ddbmetrics.provisioned-capacity.read.%v", k), *v.ProvisionedThroughput.ReadCapacityUnits, sampleRate)
			p.stats.Gauge(fmt.Sprintf("ddbmetrics.provisioned-capacity.write.%v", k), *v.ProvisionedThroughput.WriteCapacityUnits, sampleRate)

			for _, gsi := range v.GlobalSecondaryIndexes {
				p.stats.Gauge(fmt.Sprintf("ddbmetrics.provisioned-capacity.read.%v-%v", k, *gsi.IndexName), *gsi.ProvisionedThroughput.ReadCapacityUnits, sampleRate)
				p.stats.Gauge(fmt.Sprintf("ddbmetrics.provisioned-capacity.write.%v-%v", k, *gsi.IndexName), *gsi.ProvisionedThroughput.WriteCapacityUnits, sampleRate)
			}
		}
	}
}

// measureUsedCapacity reads consumed capacity metrics from the dynamo client and calculates
// values per table and index over a rolling window.
func (p *Publisher) measureUsedCapacity() {
	p.wg.Add(1)
	go func() {
		defer p.wg.Done()
		for {
			select {
			case <-p.stop:
				return
			case c := <-p.c:
				if c.consumed != nil {
					p.tables[*c.consumed.TableName] = struct{}{}
					num := p.capacityNumber(c.op, *c.consumed.TableName, nil)
					num.Increment(*c.consumed.Table.CapacityUnits)

					if c.consumed.GlobalSecondaryIndexes != nil {
						for k, v := range c.consumed.GlobalSecondaryIndexes {
							num := p.capacityNumber(c.op, *c.consumed.TableName, &k)
							num.Increment(*v.CapacityUnits)
						}
					}
				}
			}
		}
	}()
}

// recordUsedCapacity takes the calculated values from measureUsedCapacity and reports them to statsd every second.
func (p *Publisher) recordUsedCapacity() {
	p.wg.Add(1)
	go func() {
		defer p.wg.Done()
		c := time.Tick(1 * time.Second)
		for {
			select {
			case <-c:
				if p.stats != nil {
					p.usedCapacityMu.Lock()
					for k, v := range p.usedCapacity {
						avg := v.Avg(time.Now())
						op, err := strconv.Atoi(k[0:1])
						if err != nil {
							log.Print(err)
						}

						opStr := "read"
						if op == int(Write) {
							opStr = "write"
						}
						p.stats.Gauge(fmt.Sprintf("ddbmetrics.consumed-capacity.%v.%v", opStr, k[2:]), int64(avg), sampleRate)
					}
					p.usedCapacityMu.Unlock()
				}
			case <-p.stop:
				return
			}
		}
	}()
}

// fetchTableDescription queries dynamo metadata for a table
func (p *Publisher) fetchTableDescription(table string) *dynamodb.TableDescription {
	input := dynamodb.DescribeTableInput{
		TableName: aws.String(table),
	}

	out, err := p.db.DescribeTable(&input)
	if err != nil {
		log.Print(err)
		return nil
	}

	return out.Table
}

func (p *Publisher) capacityNumber(rw Operation, table string, indexName *string) *rolling.Number {
	p.usedCapacityMu.Lock()
	defer p.usedCapacityMu.Unlock()
	var capacityKey string
	if indexName == nil {
		capacityKey = fmt.Sprintf("%v-%v", rw, table)
	} else {
		capacityKey = fmt.Sprintf("%v-%v-%v", rw, table, *indexName)
	}

	_, ok := p.usedCapacity[capacityKey]
	if !ok {
		p.usedCapacity[capacityKey] = rolling.NewNumber()
	}

	return p.usedCapacity[capacityKey]
}
