package aggregate

import (
	"time"

	"github.com/netsynth/netsynth/classify"
)

// DefaultWindowMs is the default time window duration in milliseconds.
const DefaultWindowMs = 500

// Aggregate reads ClassifiedPackets from events, accumulates counts per TrafficClass
// in time windows of windowMs duration, and emits a WindowSnapshot per window.
// When done is closed, flushes the current partial window and closes the output channel.
// If onSnapshot is non-nil, it is called with each snapshot before sending to the output channel
// (used for --verbose stderr output).
func Aggregate(done <-chan struct{}, events <-chan classify.ClassifiedPacket,
	windowMs int, onSnapshot func(classify.WindowSnapshot)) <-chan classify.WindowSnapshot {

	out := make(chan classify.WindowSnapshot, 8)
	go func() {
		defer close(out)
		ticker := time.NewTicker(time.Duration(windowMs) * time.Millisecond)
		defer ticker.Stop()

		counts := make(map[classify.TrafficClass]int64)
		var totalPackets int64
		windowIndex := 0

		flush := func() {
			snap := classify.WindowSnapshot{
				Counts:       counts,
				TotalPackets: totalPackets,
				WindowIndex:  windowIndex,
			}
			if onSnapshot != nil {
				onSnapshot(snap)
			}
			out <- snap
			counts = make(map[classify.TrafficClass]int64)
			totalPackets = 0
			windowIndex++
		}

		for {
			select {
			case <-done:
				flush() // flush final partial window
				return
			case <-ticker.C:
				flush()
			case ev, ok := <-events:
				if !ok {
					flush()
					return
				}
				counts[ev.Class]++
				totalPackets++
			}
		}
	}()
	return out
}