package metrics

import (
  "time"
  "github.com/prometheus/client_golang/prometheus"
)

// Timer assumes the metrics is partitioned by one label
type Timer struct {
  startTime   map[string]time.Time
  metrics     *prometheus.HistogramVec
  measureUnit time.Duration
  labelKey    string
}

func NewTimer(metrics *prometheus.HistogramVec, unit time.Duration, labelKey string) *Timer {
  return &Timer{
    startTime:   make(map[string]time.Time),
    measureUnit: unit,
    metrics:     metrics,
    labelKey:    labelKey,
  }
}

func (i *Timer) Start(labelVal string) {
  i.startTime[labelVal] = time.Now()
}

func (i *Timer) End(labelVal string) time.Duration {
  if start, ok := i.startTime[labelVal]; ok {
    return Latency(start, time.Now())
  }
  return 0
}

func (i *Timer) Observe(measurement time.Duration, labelVal string) {
  metricsLabels := prometheus.Labels{i.labelKey: labelVal}
  i.metrics.With(metricsLabels).Observe(float64(measurement / i.measureUnit))
}

func (i *Timer) EndAndObserve(labelVal string) {
  i.Observe(i.End(labelVal), labelVal)
}

func Latency(startTime, endTime time.Time) time.Duration {
  return endTime.Sub(startTime)
}