package main

import (
	"sort"
	"sync"
	"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 {
	mu     sync.Mutex
	count  int64
	sum    time.Duration
	max    time.Duration
	bucket []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([]int64, len(latencyBounds)+1)}
}

func (h *DurationHistogram) Observe(value time.Duration) {
	if value < 0 {
		return
	}
	h.mu.Lock()
	defer h.mu.Unlock()
	h.count++
	h.sum += value
	if value > h.max {
		h.max = value
	}
	idx := sort.Search(len(latencyBounds), func(i int) bool { return value <= latencyBounds[i] })
	h.bucket[idx]++
}

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