package main

import (
	"hash/crc32"
	"sync"
	"sync/atomic"
	"time"

	"golang.org/x/sys/unix"
)

type rxStats struct {
	frames   atomic.Uint64
	bytes    atomic.Uint64
	badMagic atomic.Uint64
	badHdr   atomic.Uint64
	badLen   atomic.Uint64
	crcErr   atomic.Uint64
	rxErrs   atomic.Uint64

	// Frames counted into the interval the worker read them in, on one shared
	// host clock; read jitter is repaired at display time by the backward smear.
	newest  atomic.Int64
	buckets [rateBuckets]rxBucket
}

// One sample interval of reads per bucket, with enough of them kept that the
// smear window fits and a bucket is read long before its slot comes round
// again. smearWindow bounds how far back excess may travel and is how far
// behind real time the displayed rate runs; it must comfortably exceed the
// worst host read stall.
const (
	rateBucketNs   = int64(sampleInterval)
	rateBuckets    = 64
	rateBucketSecs = float64(rateBucketNs) / 1e9

	smearWindow = 4

	// Full wire occupancy of one bucket at line rate, in the measure gbps()
	// reports: counted bytes plus wireOverhead per frame.
	bucketWireCap = uint64(linkSpeed*1e9/8) * uint64(rateBucketNs) / 1_000_000_000
)

var rateEpochStart = time.Now()

func rateEpoch() int64 {
	return int64(time.Since(rateEpochStart)) / rateBucketNs
}

type rxBucket struct {
	epoch  atomic.Int64
	frames atomic.Uint64
	bytes  atomic.Uint64
}

// Only the owning worker writes its own buckets, so a slot coming round again is
// simply zeroed before it is claimed for the new epoch.
func (s *rxStats) commit(e int64, frames, bytes uint64) {
	b := &s.buckets[e&(rateBuckets-1)]
	if b.epoch.Load() != e {
		b.frames.Store(0)
		b.bytes.Store(0)
		b.epoch.Store(e)
	}
	b.frames.Add(frames)
	b.bytes.Add(bytes)
	if e > s.newest.Load() {
		s.newest.Store(e)
	}
}

// What this worker counted into one epoch, or nothing if that epoch has already
// fallen out of the ring.
func (s *rxStats) bucket(e int64) (frames, bytes uint64) {
	b := &s.buckets[e&(rateBuckets-1)]
	if b.epoch.Load() != e {
		return 0, 0
	}
	return b.frames.Load(), b.bytes.Load()
}

// Donate the newest bucket's excess above line rate backward into the nearest
// earlier deficits. Read jitter is purely backward — a frame is read at or
// after its arrival — so excess is frames that arrived earlier and were read
// late, and the burst drains the backlog of the stall immediately before it;
// a deficit with no later excess (a genuine wire dip) keeps its full size,
// and excess never moves forward. Frames travel with the bytes they carried,
// in the donor's proportion. Each bucket donates exactly once, on entry to
// the settled window — the donation mutates the stored values, so a donated
// byte can never display again in its donor.
func fillBack(frames, bytes []uint64) {
	i := len(frames) - 1
	for j := i - 1; j >= 0; j-- {
		wire := bytes[i] + frames[i]*wireOverhead
		if wire <= bucketWireCap {
			return
		}
		have := bytes[j] + frames[j]*wireOverhead
		if have >= bucketWireCap {
			continue
		}
		take := min(wire-bucketWireCap, bucketWireCap-have)
		mf := frames[i] * take / wire
		mb := take - mf*wireOverhead
		frames[i] -= mf
		bytes[i] -= mb
		frames[j] += mf
		bytes[j] += mb
	}
}

type rxWorker struct {
	fd     int
	batch  int
	stream uint16
	spec   *frameSpec
	stats  *rxStats
	loss   *lossWindow
	ready  *sync.WaitGroup
}

func (w *rxWorker) run(done *atomic.Bool) {
	bufs := make([][]byte, w.batch)
	for i := range bufs {
		bufs[i] = make([]byte, maxFrame)
		for j := 0; j < maxFrame; j += 4096 {
			bufs[i][j] = 0
		}
	}
	hdrs, _ := newMmsghdrs(bufs)

	w.ready.Done()

	for !done.Load() {
		n, err := recvmmsg(w.fd, hdrs, unix.MSG_WAITFORONE)
		if n <= 0 {
			if err != nil && err != unix.EAGAIN && err != unix.EINTR {
				w.stats.rxErrs.Add(1)
			}
			continue
		}
		epoch := rateEpoch()
		var frames, bytes uint64
		for i := 0; i < n; i++ {
			buf := bufs[i][:int(hdrs[i].len)]
			p, st := parseHeader(buf)
			if st != hdrOK {
				if st == hdrForeign {
					w.stats.badMagic.Add(1)
				} else {
					w.stats.badHdr.Add(1)
				}
				continue
			}
			frames++
			bytes += uint64(len(buf))

			// The ethertype this socket is bound to already says which stream the
			// frame belongs to, so a header naming another one is damaged, as is a
			// sequence number the sender never reached.
			if p.stream != w.stream || !w.loss.observe(p.seq) {
				w.stats.badHdr.Add(1)
				continue
			}

			want, ok := w.spec.expectedCRC(p.patIdx, p.payLen)
			if !ok {
				w.stats.badLen.Add(1)
				continue
			}
			pay := buf[minFrame : minFrame+p.payLen]
			if crc32.Checksum(pay, crcTable) != want {
				w.stats.crcErr.Add(1)
			}
		}
		// Nothing reads these between frames, and they share a cache line, so a
		// batch commits once rather than locking the line for every frame.
		w.stats.frames.Add(frames)
		w.stats.bytes.Add(bytes)
		if frames > 0 {
			w.stats.commit(epoch, frames, bytes)
		}
	}
}