package main

import (
	"sort"
	"sync/atomic"
	"time"
)

var latencyBounds = []time.Duration{
	time.Millisecond, 3 * time.Millisecond, 5 * time.Millisecond, 10 * time.Millisecond,
	25 * time.Millisecond, 50 * time.Millisecond, 100 * time.Millisecond, 200 * time.Millisecond,
	500 * time.Millisecond, time.Second, 2 * time.Second, 5 * time.Second, 10 * time.Second,
}

type DurationHistogram struct {
	count  atomic.Int64
	sum    atomic.Int64
	max    atomic.Int64
	bucket []atomic.Int64
}

type HistogramSnapshot struct {
	Count   int64            `json:"count"`
	SumMS   float64          `json:"sum_ms"`
	MaxMS   float64          `json:"max_ms"`
	Buckets map[string]int64 `json:"buckets"`
}

func NewDurationHistogram() *DurationHistogram {
	return &DurationHistogram{bucket: make([]atomic.Int64, len(latencyBounds)+1)}
}

func (h *DurationHistogram) Observe(value time.Duration) {
	if value < 0 {
		return
	}
	nanos := int64(value)
	h.count.Add(1)
	h.sum.Add(nanos)
	for previous := h.max.Load(); nanos > previous; previous = h.max.Load() {
		if h.max.CompareAndSwap(previous, nanos) {
			break
		}
	}
	idx := sort.Search(len(latencyBounds), func(i int) bool { return value <= latencyBounds[i] })
	h.bucket[idx].Add(1)
}

func (h *DurationHistogram) Snapshot() HistogramSnapshot {
	buckets := make(map[string]int64, len(h.bucket))
	for i := range h.bucket {
		label := "+Inf"
		if i < len(latencyBounds) {
			label = latencyBounds[i].String()
		}
		buckets[label] = h.bucket[i].Load()
	}
	return HistogramSnapshot{
		Count: h.count.Load(), SumMS: float64(h.sum.Load()) / float64(time.Millisecond),
		MaxMS: float64(h.max.Load()) / float64(time.Millisecond), Buckets: buckets,
	}
}