package utils

import (
	"sync/atomic"
	"time"
)

type RollingStats struct {
	buckets     []atomic.Int64
	bucketIndex atomic.Int32
	numBuckets  int32
	bucketSize  time.Duration
}

func NewRollingStats(numBuckets int32, bucketSize time.Duration) *RollingStats {
	rs := &RollingStats{
		buckets:    make([]atomic.Int64, numBuckets),
		numBuckets: numBuckets,
		bucketSize: bucketSize,
	}
	go rs.rotateBuckets()
	return rs
}

func (rs *RollingStats) rotateBuckets() {
	ticker := time.NewTicker(rs.bucketSize)
	defer ticker.Stop()
	for range ticker.C {
		newIdx := (rs.bucketIndex.Load() + 1) % int32(rs.numBuckets)
		rs.bucketIndex.Store(newIdx)
		rs.buckets[newIdx].Store(0)
	}
}

func (rs *RollingStats) Increment() {
	idx := rs.bucketIndex.Load()
	rs.buckets[idx].Add(1)
}

func (rs *RollingStats) GetRPS() float64 {
	var total int64
	for i := 0; i < int(rs.numBuckets); i++ {
		total += rs.buckets[i].Load()
	}

	windowSeconds := float64(rs.numBuckets) * rs.bucketSize.Seconds()
	return float64(total) / windowSeconds
}

func (rs *RollingStats) GetTotal() int64 {
	var total int64
	for i := 0; i < int(rs.numBuckets); i++ {
		total += rs.buckets[i].Load()
	}
	return total
}