package statrecorder

import (
	"strconv"
	"strings"
	"sync"
	"time"

	"github.com/cactus/go-statsd-client/statsd"
)

// StatRecorder implements
// github.com/cactus/go-statsd-client/statsd's Statter, SubStatter,
// and StatSender interfaces by recording statistics in memory. This
// is useful for tests: one can instantiate a StatRecorder and then
// make assertions about the statistics it has received.
//
// StatRecorder never samples data, regardless of any sampling rates
// passed in to its functions. It also never 'flushes' in any sense,
// so counters will continue accumulating indefinitely. Create a new
// StatRecorder per test to guarantee clean statistics.
//
// StatRecorder is safe for concurrent usage.
type StatRecorder struct {
	// nameLock protects access to the recorder's metric naming
	nameLock sync.Mutex
	prefix   string

	// dataLock protects access to the recorder's data
	dataLock sync.Mutex
	counters map[string]int
	gauges   map[string]int
	timers   map[string][]time.Duration
	sets     map[string]map[string]struct{}
	raws     map[string][]string
}

// Compiler assertions that we meet the interfaces in statsd
var _ statsd.Statter = &StatRecorder{}
var _ statsd.SubStatter = &StatRecorder{}
var _ statsd.StatSender = &StatRecorder{}

// New creates a StatRecorder that will use the given prefix in front
// of all metrics it tracks.
func New(prefix string) *StatRecorder {
	return &StatRecorder{
		prefix:   prefix,
		counters: make(map[string]int),
		gauges:   make(map[string]int),
		timers:   make(map[string][]time.Duration),
		sets:     make(map[string]map[string]struct{}),
		raws:     make(map[string][]string),
	}
}

// MetricName generates the proper metric name, prepending m with
// prefix. It's copied from unexported code in
// github.com/cactus/go-statsd-client/statsd.
func (sr *StatRecorder) MetricName(m string) string {
	sr.nameLock.Lock()
	defer sr.nameLock.Unlock()
	m = strings.TrimLeft(m, ".")
	if sr.prefix != "" && m != "" {
		return sr.prefix + "." + m
	}
	return sr.prefix + m
}

// NewSubStatter creates a SubStatter from sr, appending s to sr's
// prefix for the SubStatter.
func (sr *StatRecorder) NewSubStatter(s string) statsd.SubStatter {
	return New(sr.MetricName(s))
}

// SetPrefix changes sr's prefix.
func (sr *StatRecorder) SetPrefix(p string) {
	sr.nameLock.Lock()
	defer sr.nameLock.Unlock()
	sr.prefix = p
}

// Close does nothing and always returns nil. It exists to fulfill interfaces.
func (sr *StatRecorder) Close() error { return nil }

// SetSamplerFunc does nothing. It exists to fulfill interfaces.
func (sr *StatRecorder) SetSamplerFunc(statsd.SamplerFunc) {}

// Inc increments a counter. It always returns nil. The sampling rate
// s is ignored.
func (sr *StatRecorder) Inc(m string, v int64, s float32) error {
	sr.dataLock.Lock()
	defer sr.dataLock.Unlock()
	sr.counters[sr.MetricName(m)] += int(v)
	return nil
}

// Dec decrements a counter. It always returns nil. The sampling rate
// s is ignored.
func (sr *StatRecorder) Dec(m string, v int64, s float32) error {
	return sr.Inc(m, -v, s)
}

// Gauge sets a gauge to a value. It always returns nil. The sampling
// rate s is ignored.
func (sr *StatRecorder) Gauge(m string, v int64, s float32) error {
	sr.dataLock.Lock()
	defer sr.dataLock.Unlock()
	sr.gauges[sr.MetricName(m)] = int(v)
	return nil
}

// GaugeDelta changes a gauge by a value. It always returns nil. The sampling
// rate s is ignored.
func (sr *StatRecorder) GaugeDelta(m string, v int64, s float32) error {
	sr.dataLock.Lock()
	defer sr.dataLock.Unlock()
	sr.gauges[sr.MetricName(m)] += int(v)
	return nil
}

// Timing records timing information. The value v is interpreted as
// milliseconds. It always returns nil. The sampling rate s is
// ignored.
func (sr *StatRecorder) Timing(m string, v int64, s float32) error {
	sr.dataLock.Lock()
	defer sr.dataLock.Unlock()
	return sr.TimingDuration(m, time.Duration(v)*time.Millisecond, s)
}

// TimingDuration records timing information. It always returns
// nil. The sampling rate s is ignored.
func (sr *StatRecorder) TimingDuration(m string, v time.Duration, s float32) error {
	sr.dataLock.Lock()
	defer sr.dataLock.Unlock()
	sr.timers[sr.MetricName(m)] = append(sr.timers[sr.MetricName(m)], v)
	return nil
}

// Set adds v to a set of strings. It always returns nil. The sampling
// rate s is ignored.
func (sr *StatRecorder) Set(m string, v string, s float32) error {
	sr.dataLock.Lock()
	defer sr.dataLock.Unlock()
	set, ok := sr.sets[sr.MetricName(m)]
	if !ok {
		set = make(map[string]struct{})
		sr.sets[sr.MetricName(m)] = set
	}
	set[v] = struct{}{}
	return nil
}

// SetInt converts v to a string and adds it to a set of strings. It
// always returns nil. The sampling rate s is ignored.
func (sr *StatRecorder) SetInt(m string, v int64, s float32) error {
	sr.dataLock.Lock()
	defer sr.dataLock.Unlock()
	return sr.Set(m, strconv.FormatInt(v, 10), s)
}

// Raw appends v to a list of raw metrics that sr has received. It
// always returns nil. The sampling rate s is ignored.
func (sr *StatRecorder) Raw(m string, v string, s float32) error {
	sr.dataLock.Lock()
	defer sr.dataLock.Unlock()
	sr.raws[sr.MetricName(m)] = append(sr.raws[sr.MetricName(m)], v)
	return nil
}

// CounterValue is a convenience method which retrieves the counter
// specified by the metric m, prepending m with the appropriate
// prefix.
func (sr *StatRecorder) CounterValue(m string) int {
	sr.dataLock.Lock()
	defer sr.dataLock.Unlock()
	return sr.counters[sr.MetricName(m)]
}

// AllCounters returns a copy of the complete set of counters as a map
// of metric name to current value. The StatRecorder's prefix is
// included in the metric names.
func (sr *StatRecorder) AllCounters() map[string]int {
	sr.dataLock.Lock()
	defer sr.dataLock.Unlock()
	copy := make(map[string]int, len(sr.counters))
	for k, v := range sr.counters {
		copy[k] = v
	}
	return copy
}

// GaugeValue is a convenience method which retrieves the gauge
// specified by the metric m, prepending m with the appropriate
// prefix.
func (sr *StatRecorder) GaugeValue(m string) int {
	sr.dataLock.Lock()
	defer sr.dataLock.Unlock()
	return sr.gauges[sr.MetricName(m)]
}

// AllGauges returns a copy of the complete set of gauges as a map
// of metric name to current value. The StatRecorder's prefix is
// included in the metric names.
func (sr *StatRecorder) AllGauges() map[string]int {
	sr.dataLock.Lock()
	defer sr.dataLock.Unlock()
	copy := make(map[string]int, len(sr.gauges))
	for k, v := range sr.gauges {
		copy[k] = v
	}
	return copy
}

// SetValue is a convenience method which retrieves the set specified
// by the metric m, prepending m with the appropriate prefix.
func (sr *StatRecorder) SetValue(m string) map[string]struct{} {
	sr.dataLock.Lock()
	defer sr.dataLock.Unlock()
	// Copy the set so we don't return a mutable value here
	return copySet(sr.sets[sr.MetricName(m)])
}

func copySet(s map[string]struct{}) map[string]struct{} {
	s2 := make(map[string]struct{})
	for k := range s {
		s2[k] = struct{}{}
	}
	return s2
}

// AllSets returns a copy of the complete set of Sets as a map
// of metric name to current value. The StatRecorder's prefix is
// included in the metric names.
func (sr *StatRecorder) AllSets() map[string]map[string]struct{} {
	sr.dataLock.Lock()
	defer sr.dataLock.Unlock()
	copy := make(map[string]map[string]struct{}, len(sr.gauges))
	for k, v := range sr.sets {
		copy[k] = copySet(v)
	}
	return copy
}

func copyTimers(ts []time.Duration) []time.Duration {
	ts2 := make([]time.Duration, len(ts))
	copy(ts2, ts)
	return ts2
}

// TimingValues is a convenience method which retrieves the timer
// values which have been recorded using the metric m, prepending m
// with the appropriate prefix.
func (sr *StatRecorder) TimingValues(m string) []time.Duration {
	sr.dataLock.Lock()
	defer sr.dataLock.Unlock()
	// Copy the timers so we don't return a mutable value here
	return copyTimers(sr.timers[sr.MetricName(m)])
}

// AllTimings returns a copy of the complete set of timers as a map of
// metric name to current value. The StatRecorder's prefix is included
// in the metric names.
func (sr *StatRecorder) AllTimings() map[string][]time.Duration {
	sr.dataLock.Lock()
	defer sr.dataLock.Unlock()
	copiedTimers := make(map[string][]time.Duration, len(sr.gauges))
	for k, v := range sr.timers {
		copiedTimers[k] = copyTimers(v)
	}
	return copiedTimers
}
