package main

import (
	"math"
	"sync"
)

const statRingSize = 512

// SourceStat accumulates per-source UDP performance metrics.
// Thread-safe; RecordFragment is called from the UDPClient goroutine.
type SourceStat struct {
	mu sync.Mutex

	seenFirst   bool
	lastCounter uint32
	TotalRx     uint64
	TotalLost   uint64

	// Cycle-time ring (seconds between consecutive DATA completions)
	ctRing   [statRingSize]float64
	ctHead   int
	ctFull   bool
	lastRxNs int64

	// Per-cycle accumulators (reset after each DATA completion)
	fragCount int
	byteCount int

	fragRing [statRingSize]int
	byteRing [statRingSize]int
}

// RecordFragment is called for every UDP datagram of a DATA packet.
//   complete: this fragment completed the DATA reassembly.
//   nBytes:   raw datagram size (header+payload).
func (s *SourceStat) RecordFragment(counter uint32, nBytes int, arrivalNs int64, complete bool) {
	s.mu.Lock()
	defer s.mu.Unlock()

	s.fragCount++
	s.byteCount += nBytes

	if !complete {
		return
	}

	s.TotalRx++
	if s.seenFirst {
		if delta := counter - s.lastCounter; delta > 1 {
			s.TotalLost += uint64(delta - 1)
		}
	} else {
		s.seenFirst = true
	}
	s.lastCounter = counter

	if s.lastRxNs != 0 {
		idx := s.ctHead
		s.ctRing[idx] = float64(arrivalNs-s.lastRxNs) * 1e-9
		s.fragRing[idx] = s.fragCount
		s.byteRing[idx] = s.byteCount
		s.ctHead = (s.ctHead + 1) % statRingSize
		if s.ctHead == 0 {
			s.ctFull = true
		}
	}
	s.lastRxNs = arrivalNs
	s.fragCount = 0
	s.byteCount = 0
}

// Snapshot computes and returns a StatInfo for broadcast.
func (s *SourceStat) Snapshot() StatInfo {
	s.mu.Lock()
	defer s.mu.Unlock()

	n := s.ctHead
	if s.ctFull {
		n = statRingSize
	}

	si := StatInfo{TotalReceived: s.TotalRx, TotalLost: s.TotalLost}
	if n == 0 {
		return si
	}

	sum, sumSq := 0.0, 0.0
	minV, maxV := math.MaxFloat64, 0.0
	fragSum, byteSum := 0, 0
	for i := 0; i < n; i++ {
		v := s.ctRing[i]
		sum += v
		sumSq += v * v
		if v < minV {
			minV = v
		}
		if v > maxV {
			maxV = v
		}
		fragSum += s.fragRing[i]
		byteSum += s.byteRing[i]
	}
	avg := sum / float64(n)
	variance := sumSq/float64(n) - avg*avg
	if variance < 0 {
		variance = 0
	}
	stdv := math.Sqrt(variance)

	si.CycleAvgMs = avg * 1e3
	si.CycleStdMs = stdv * 1e3
	si.CycleMinMs = minV * 1e3
	si.CycleMaxMs = maxV * 1e3
	si.RateHz = 1.0 / avg
	si.RateStdHz = stdv / (avg * avg)
	si.FragsPerCycle = float64(fragSum) / float64(n)
	si.BytesPerCycle = float64(byteSum) / float64(n)

	const nBins = 20
	si.CycleHistMin = minV * 1e3
	si.CycleHistMax = maxV * 1e3
	si.CycleHist = make([]int, nBins)
	span := maxV - minV
	for i := 0; i < n; i++ {
		var bin int
		if span > 0 {
			bin = int((s.ctRing[i] - minV) / span * float64(nBins))
			if bin >= nBins {
				bin = nBins - 1
			}
		} else {
			bin = nBins / 2
		}
		si.CycleHist[bin]++
	}
	return si
}

// StatInfo is the JSON snapshot for one source sent to the frontend.
type StatInfo struct {
	TotalReceived uint64  `json:"totalReceived"`
	TotalLost     uint64  `json:"totalLost"`
	RateHz        float64 `json:"rateHz"`
	RateStdHz     float64 `json:"rateStdHz"`
	FragsPerCycle float64 `json:"fragsPerCycle"`
	BytesPerCycle float64 `json:"bytesPerCycle"`
	CycleAvgMs    float64 `json:"cycleAvgMs"`
	CycleStdMs    float64 `json:"cycleStdMs"`
	CycleMinMs    float64 `json:"cycleMinMs"`
	CycleMaxMs    float64 `json:"cycleMaxMs"`
	CycleHist     []int   `json:"cycleHist"`
	CycleHistMin  float64 `json:"cycleHistMin"`
	CycleHistMax  float64 `json:"cycleHistMax"`
}