package main

import (
	"math/bits"
	"sync"
	"sync/atomic"
)

const (
	lossSlots = 1 << 16
	lossWords = lossSlots / 64
)

type lossWindow struct {
	mu     sync.Mutex
	base   uint64
	inited bool
	bits   []uint64
	lost   atomic.Uint64
	late   atomic.Uint64

	// The sending half of this stream. Both halves are in this process and the
	// receiving socket ignores its own outgoing frames, so this window is fed by
	// that sender alone and by nothing else on the wire.
	sent *atomic.Uint64
}

// Takes the sending halves rather than a count, so a window cannot be built
// without the bound it judges sequence numbers against.
func newLossWindows(tx []*txStats) []lossWindow {
	w := make([]lossWindow, len(tx))
	for i := range w {
		w[i].bits = make([]uint64, lossWords)
		w[i].sent = &tx[i].sent
	}
	return w
}

// A sequence number is only judged once it falls out of the window, so frames
// still queued in another rx worker are never miscounted as lost.
//
// Reports whether the sequence number could have come off the wire at all. One
// above what the sender has reached is a damaged header rather than a gap, and
// is refused before it can touch the base. Believing one would drag the window
// up past every real sequence, write off the span in between as lost, and leave
// every frame after it arriving below the base and counted late for the rest of
// the run. Nothing that arrives later is evidence enough to undo that, which is
// why the sender's own frontier is the guard rather than a limit on how far a
// single step may move.
func (w *lossWindow) observe(seq uint64) bool {
	if seq >= w.sent.Load() {
		return false
	}
	w.mu.Lock()
	if !w.inited {
		// Start half a window below the first sequence seen, so frames another
		// rx worker is still holding land inside the window rather than late.
		if seq > lossSlots/2 {
			w.base = seq - lossSlots/2
		}
		w.inited = true
	}
	if seq < w.base {
		w.mu.Unlock()
		w.late.Add(1)
		return true
	}
	if seq >= w.base+lossSlots {
		w.evict(seq - lossSlots + 1)
	}
	idx := seq & (lossSlots - 1)
	w.bits[idx>>6] |= 1 << (idx & 63)
	w.mu.Unlock()
	return true
}

func (w *lossWindow) evict(newBase uint64) {
	span := newBase - w.base
	if span >= lossSlots {
		var missing uint64
		for i := range w.bits {
			missing += uint64(64 - bits.OnesCount64(w.bits[i]))
			w.bits[i] = 0
		}
		w.lost.Add(missing + span - lossSlots)
		w.base = newBase
		return
	}
	var missing uint64
	for s := w.base; s < newBase; s++ {
		idx := s & (lossSlots - 1)
		word, bit := idx>>6, uint64(1)<<(idx&63)
		if w.bits[word]&bit == 0 {
			missing++
			continue
		}
		w.bits[word] &^= bit
	}
	w.lost.Add(missing)
	w.base = newBase
}