package metrics

import (
	"strings"
	"sync"

	"github.com/prometheus/client_golang/prometheus"
)

// GaugeVec of prometheus has no way to delete or reset based on partial labels
// See https://github.com/prometheus/client_golang/issues/834
// So a custom GaugeVec is implemented to set labels easily and atomically.
type GaugeVec struct {
	lock sync.RWMutex
	prom *prometheus.GaugeVec

	keyToLabels map[string]prometheus.Labels
	keyNames    []string
}

func NewGaugeVec(opts prometheus.GaugeOpts, labelNames []string, keyNames []string) *GaugeVec {
	return &GaugeVec{
		prom:        prometheus.NewGaugeVec(opts, labelNames),
		keyToLabels: make(map[string]prometheus.Labels),
		keyNames:    keyNames,
	}
}

func (v *GaugeVec) Set(labels prometheus.Labels, val float64) {
	keyVals := []string{}
	for _, k := range v.keyNames {
		keyVals = append(keyVals, labels[k])
	}
	key := strings.Join(keyVals, "/")

	v.lock.Lock()
	defer v.lock.Unlock()

	oldLabels, ok := v.keyToLabels[key]
	if ok {
		v.prom.Delete(oldLabels)
	}
	v.keyToLabels[key] = labels
	g := v.prom.With(labels)
	g.Set(val)
}

func (v *GaugeVec) Delete(labels prometheus.Labels) {
	keyVals := []string{}
	for _, k := range v.keyNames {
		keyVals = append(keyVals, labels[k])
	}
	key := strings.Join(keyVals, "/")

	v.lock.Lock()
	defer v.lock.Unlock()

	oldLabels, ok := v.keyToLabels[key]
	if ok {
		v.prom.Delete(oldLabels)
		delete(v.keyToLabels, key)
	}
}

func (v *GaugeVec) Collect(ch chan<- prometheus.Metric) {
	v.lock.RLock()
	defer v.lock.RUnlock()
	v.prom.Collect(ch)
}

func (v *GaugeVec) Describe(ch chan<- *prometheus.Desc) {
	v.prom.Describe(ch)
}