package main

import (
	"flag"
	"fmt"
	"net"
	"os"
	"os/signal"
	"strconv"
	"strings"
	"sync"
	"sync/atomic"
	"syscall"
	"time"

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

const (
	fanoutHash     = 0
	fanoutLB       = 1
	fanoutCPU      = 2
	fanoutRollover = 3
)

const wireOverhead = 24

type endpoint struct {
	name string
	idx  int
	mac  [6]byte
	mtu  int
}

type direction struct {
	label   string
	tx      endpoint
	rx      endpoint
	spec    *frameSpec
	txStats []*txStats
	rxStats []*rxStats
	streams []streamState
	txFDs   []int
	rxFDs   []int
	reports chan string

	prev  sample
	drops uint64
	nicTX nicCounters
	nicRX nicCounters
}

type sample struct {
	txFrames, txBytes uint64
	rxFrames, rxBytes uint64
	expected, count   uint64
	crcErr, badMagic  uint64
	badLen            uint64
	txErrs, txShort   uint64
}

func lookupEndpoint(name string) (endpoint, error) {
	ifi, err := net.InterfaceByName(name)
	if err != nil {
		return endpoint{}, err
	}
	if len(ifi.HardwareAddr) != 6 {
		return endpoint{}, fmt.Errorf("%s: expected 6-byte MAC, got %q", name, ifi.HardwareAddr)
	}
	var mac [6]byte
	copy(mac[:], ifi.HardwareAddr)
	return endpoint{name: name, idx: ifi.Index, mac: mac, mtu: ifi.MTU}, nil
}

func parseSizes(s string) ([]int, error) {
	var out []int
	for _, f := range strings.Split(s, ",") {
		f = strings.TrimSpace(f)
		if f == "" {
			continue
		}
		v, err := strconv.Atoi(f)
		if err != nil {
			return nil, fmt.Errorf("bad size %q: %w", f, err)
		}
		if v < minFrame {
			return nil, fmt.Errorf("size %d below minimum %d", v, minFrame)
		}
		out = append(out, v)
	}
	if len(out) == 0 {
		return nil, fmt.Errorf("no sizes given")
	}
	return out, nil
}

func parseCPUs(s string) ([]int, error) {
	if strings.TrimSpace(s) == "" {
		return nil, nil
	}
	var out []int
	for _, f := range strings.Split(s, ",") {
		v, err := strconv.Atoi(strings.TrimSpace(f))
		if err != nil {
			return nil, fmt.Errorf("bad cpu %q: %w", f, err)
		}
		out = append(out, v)
	}
	return out, nil
}

func cpuAt(cpus []int, i int) int {
	if len(cpus) == 0 {
		return -1
	}
	return cpus[i%len(cpus)]
}

func fanoutMode(name string, nsock int) (int, error) {
	if nsock < 2 {
		return -1, nil
	}
	switch name {
	case "none":
		return -1, nil
	case "hash":
		return fanoutHash, nil
	case "lb":
		return fanoutLB, nil
	case "cpu":
		return fanoutCPU, nil
	case "rollover":
		return fanoutRollover, nil
	}
	return 0, fmt.Errorf("unknown fanout mode %q", name)
}

func (d *direction) snapshot() sample {
	var s sample
	for _, t := range d.txStats {
		s.txFrames += t.frames.Load()
		s.txBytes += t.bytes.Load()
		s.txErrs += t.errs.Load()
		s.txShort += t.short.Load()
	}
	for _, r := range d.rxStats {
		s.rxFrames += r.frames.Load()
		s.rxBytes += r.bytes.Load()
		s.crcErr += r.crcErr.Load()
		s.badMagic += r.badMagic.Load()
		s.badLen += r.badLen.Load()
	}
	for i := range d.streams {
		c := d.streams[i].count.Load()
		if c == 0 {
			continue
		}
		s.count += c
		s.expected += d.streams[i].maxSeq.Load() + 1
	}
	return s
}

func (d *direction) sampleDrops() {
	for _, fd := range d.rxFDs {
		d.drops += packetDrops(fd)
	}
}

func gbps(bytes, frames uint64, secs float64) float64 {
	return float64((bytes+frames*wireOverhead)*8) / secs / 1e9
}

func (d *direction) reportInterval(secs float64) string {
	now := d.snapshot()
	p := d.prev
	d.prev = now

	txF := now.txFrames - p.txFrames
	txB := now.txBytes - p.txBytes
	rxF := now.rxFrames - p.rxFrames
	rxB := now.rxBytes - p.rxBytes

	lost := int64(now.expected) - int64(now.count)

	before := d.drops
	d.sampleDrops()
	drops := d.drops - before

	return fmt.Sprintf("%s tx %8.0f pps %6.2f Gb/s | rx %8.0f pps %6.2f Gb/s | lost(cum) %6d crc %d badmagic %d kdrop %d txshort %d txerr %d",
		d.label,
		float64(txF)/secs, gbps(txB, txF, secs),
		float64(rxF)/secs, gbps(rxB, rxF, secs),
		lost,
		now.crcErr-p.crcErr,
		now.badMagic-p.badMagic,
		drops,
		now.txShort-p.txShort,
		now.txErrs-p.txErrs,
	)
}

func (d *direction) reportNIC() string {
	tx := readNIC(d.tx.name)
	rx := readNIC(d.rx.name)
	var parts []string
	if s := tx.diff(d.nicTX); s != "" {
		parts = append(parts, "  "+d.tx.name+"(tx): "+s)
	}
	if s := rx.diff(d.nicRX); s != "" {
		parts = append(parts, "  "+d.rx.name+"(rx): "+s)
	}
	d.nicTX = tx
	d.nicRX = rx
	return strings.Join(parts, "\n")
}

func buildDirection(label string, tx, rx endpoint, patIdx int, sizes []int, cfg config) (*direction, error) {
	d := &direction{
		label:   label,
		tx:      tx,
		rx:      rx,
		spec:    newFrameSpec(patIdx, rx.mac, tx.mac, sizes),
		streams: make([]streamState, cfg.txWorkers),
		reports: make(chan string, 64),
	}
	d.nicTX = readNIC(tx.name)
	d.nicRX = readNIC(rx.name)

	fm, err := fanoutMode(cfg.fanout, cfg.rxWorkers)
	if err != nil {
		return nil, err
	}

	for i := 0; i < cfg.txWorkers; i++ {
		fd, err := openTxSocket(tx.idx, cfg.sndbuf)
		if err != nil {
			return nil, fmt.Errorf("%s tx socket: %w", label, err)
		}
		d.txFDs = append(d.txFDs, fd)
		d.txStats = append(d.txStats, &txStats{})
	}
	for i := 0; i < cfg.rxWorkers; i++ {
		fd, err := openRxSocket(rx.idx, cfg.rcvbuf, cfg.fanoutID, fm)
		if err != nil {
			return nil, fmt.Errorf("%s rx socket: %w", label, err)
		}
		d.rxFDs = append(d.rxFDs, fd)
		d.rxStats = append(d.rxStats, &rxStats{})
	}
	return d, nil
}

func (d *direction) start(wg *sync.WaitGroup, doneTx, doneRx *atomic.Bool, cfg config, rxReady *sync.WaitGroup, startTx <-chan struct{}) {
	for i, fd := range d.txFDs {
		w := &txWorker{
			fd:      fd,
			stream:  uint16(i),
			spec:    d.spec,
			batch:   cfg.batch,
			cpu:     cpuAt(cfg.txCPUs, i),
			stats:   d.txStats[i],
			startTx: startTx,
		}
		wg.Add(1)
		go func() {
			defer wg.Done()
			w.run(doneTx)
		}()
	}
	for i, fd := range d.rxFDs {
		w := &rxWorker{
			fd:      fd,
			batch:   cfg.batch,
			cpu:     cpuAt(cfg.rxCPUs, i),
			verify:  cfg.verify,
			spec:    d.spec,
			stats:   d.rxStats[i],
			streams: d.streams,
			reports: d.reports,
			ready:   rxReady,
		}
		wg.Add(1)
		go func() {
			defer wg.Done()
			w.run(doneRx)
		}()
	}
}

func (d *direction) close() {
	for _, fd := range d.txFDs {
		unix.Close(fd)
	}
	for _, fd := range d.rxFDs {
		unix.Close(fd)
	}
}

type config struct {
	txWorkers int
	rxWorkers int
	batch     int
	verify    bool
	sndbuf    int
	rcvbuf    int
	fanout    string
	fanoutID  int
	txCPUs    []int
	rxCPUs    []int
}

func main() {
	var (
		aName    = flag.String("a", "", "first interface")
		bName    = flag.String("b", "", "second interface")
		duration = flag.Duration("duration", 0, "run time, 0 for until interrupted")
		sizesArg = flag.String("sizes", "64,128,256,512,1024,1280,1514", "frame sizes in bytes, excluding FCS, cycled per packet")
		patArg   = flag.String("pattern", "prbs", "payload pattern")
		txN      = flag.Int("tx", 4, "tx workers per direction")
		rxN      = flag.Int("rx", 4, "rx workers per direction")
		batch    = flag.Int("batch", 64, "frames per sendmmsg/recvmmsg call")
		verify   = flag.Bool("verify", true, "verify payload CRC32C on receive")
		interval = flag.Duration("interval", time.Second, "report interval")
		drain    = flag.Duration("drain", 500*time.Millisecond, "keep receiving this long after tx stops, so in-flight frames are not counted as lost")
		fanout   = flag.String("fanout", "lb", "rx fanout mode: none, hash, lb, cpu, rollover")
		duplex   = flag.Bool("duplex", true, "run both directions simultaneously")
		sndbuf   = flag.Int("sndbuf", 8<<20, "SO_SNDBUFFORCE per tx socket")
		rcvbuf   = flag.Int("rcvbuf", 64<<20, "SO_RCVBUFFORCE per rx socket")
		txCPUs   = flag.String("txcpus", "", "comma-separated CPUs to pin tx workers to")
		rxCPUs   = flag.String("rxcpus", "", "comma-separated CPUs to pin rx workers to")
		tune     = flag.Bool("tune", true, "check host settings at startup and correct them; without this they are only reported")
		governor = flag.String("governor", "performance", "required cpufreq governor")
		coalesce = flag.Uint("coalesce-usecs", 25, "required fixed rx/tx coalesce usecs, with adaptive coalescing off")
		rxRing   = flag.Uint("rx-ring", 8160, "required rx ring size, clamped to hardware maximum")
		txRing   = flag.Uint("tx-ring", 4096, "required tx ring size, clamped to hardware maximum")
		setRings = flag.Bool("set-rings", true, "allow ring resizing at startup, which resets the link")
	)
	flag.Parse()

	tcfg := tuneConfig{
		enabled:  *tune,
		governor: *governor,
		rxUsecs:  uint32(*coalesce),
		txUsecs:  uint32(*coalesce),
		rxRing:   uint32(*rxRing),
		txRing:   uint32(*txRing),
		setRings: *setRings,
	}

	if err := run(*aName, *bName, *sizesArg, *patArg, *fanout, *txCPUs, *rxCPUs,
		*duration, *interval, *drain, *txN, *rxN, *batch, *sndbuf, *rcvbuf, *verify, *duplex, tcfg); err != nil {
		fmt.Fprintln(os.Stderr, "error:", err)
		os.Exit(1)
	}
}

func run(aName, bName, sizesArg, patArg, fanout, txCPUsArg, rxCPUsArg string,
	duration, interval, drain time.Duration, txN, rxN, batch, sndbuf, rcvbuf int,
	verify, duplex bool, tcfg tuneConfig) error {

	if aName == "" || bName == "" {
		return fmt.Errorf("both -a and -b are required")
	}
	sizes, err := parseSizes(sizesArg)
	if err != nil {
		return err
	}
	patIdx, err := patternIndex(patArg)
	if err != nil {
		return err
	}
	txCPUs, err := parseCPUs(txCPUsArg)
	if err != nil {
		return err
	}
	rxCPUs, err := parseCPUs(rxCPUsArg)
	if err != nil {
		return err
	}
	a, err := lookupEndpoint(aName)
	if err != nil {
		return err
	}
	b, err := lookupEndpoint(bName)
	if err != nil {
		return err
	}
	for _, e := range []endpoint{a, b} {
		for _, s := range sizes {
			if s > e.mtu+ethHdrLen {
				return fmt.Errorf("size %d exceeds %s MTU %d (max frame %d)", s, e.name, e.mtu, e.mtu+ethHdrLen)
			}
		}
	}

	ifnames := []string{a.name, b.name}
	fmt.Println("host settings:")
	var fatal []string
	for _, r := range applyTuning(tcfg, ifnames) {
		fmt.Println(r)
		if r.fatal {
			fatal = append(fatal, r.item)
		}
	}
	if len(fatal) > 0 {
		return fmt.Errorf("cannot test with %s in this state", strings.Join(fatal, ", "))
	}

	cfg := config{
		txWorkers: txN,
		rxWorkers: rxN,
		batch:     batch,
		verify:    verify,
		sndbuf:    sndbuf,
		rcvbuf:    rcvbuf,
		fanout:    fanout,
		txCPUs:    txCPUs,
		rxCPUs:    rxCPUs,
	}

	var dirs []*direction
	cfgA := cfg
	cfgA.fanoutID = 0x4341
	d0, err := buildDirection(a.name+"->"+b.name, a, b, patIdx, sizes, cfgA)
	if err != nil {
		return err
	}
	dirs = append(dirs, d0)
	if duplex {
		cfgB := cfg
		cfgB.fanoutID = 0x4342
		d1, err := buildDirection(b.name+"->"+a.name, b, a, patIdx, sizes, cfgB)
		if err != nil {
			return err
		}
		dirs = append(dirs, d1)
	}
	defer func() {
		for _, d := range dirs {
			d.close()
		}
	}()

	fmt.Printf("pattern=%s sizes=%v tx=%d rx=%d batch=%d verify=%v fanout=%s duplex=%v\n",
		patterns[patIdx].name, sizes, txN, rxN, batch, verify, fanout, duplex)
	for _, d := range dirs {
		fmt.Printf("  %s  %02x:%02x:%02x:%02x:%02x:%02x -> %02x:%02x:%02x:%02x:%02x:%02x  ethertype 0x%04x\n",
			d.label,
			d.tx.mac[0], d.tx.mac[1], d.tx.mac[2], d.tx.mac[3], d.tx.mac[4], d.tx.mac[5],
			d.rx.mac[0], d.rx.mac[1], d.rx.mac[2], d.rx.mac[3], d.rx.mac[4], d.rx.mac[5],
			etherType)
	}
	fmt.Printf("  sndbuf %d rcvbuf %d (as granted by kernel)\n",
		sockBufSize(dirs[0].txFDs[0], unix.SO_SNDBUF),
		sockBufSize(dirs[0].rxFDs[0], unix.SO_RCVBUF))

	var doneTx, doneRx atomic.Bool
	var wg sync.WaitGroup
	var rxReady sync.WaitGroup
	startTx := make(chan struct{})
	for _, d := range dirs {
		rxReady.Add(len(d.rxFDs))
	}
	for _, d := range dirs {
		d.start(&wg, &doneTx, &doneRx, cfg, &rxReady, startTx)
	}
	rxReady.Wait()

	sig := make(chan os.Signal, 1)
	signal.Notify(sig, syscall.SIGINT, syscall.SIGTERM)

	start := time.Now()
	close(startTx)
	tick := time.NewTicker(interval)
	defer tick.Stop()

	var deadline <-chan time.Time
	if duration > 0 {
		t := time.NewTimer(duration)
		defer t.Stop()
		deadline = t.C
	}

	last := time.Now()
loop:
	for {
		select {
		case <-sig:
			break loop
		case <-deadline:
			break loop
		case now := <-tick.C:
			secs := now.Sub(last).Seconds()
			last = now
			for _, w := range verifyTuning(tcfg, ifnames) {
				fmt.Println(w)
			}
			for _, d := range dirs {
				fmt.Println(d.reportInterval(secs))
				if s := d.reportNIC(); s != "" {
					fmt.Println(s)
				}
				for {
					select {
					case msg := <-d.reports:
						fmt.Println("  " + d.label + ": " + msg)
						continue
					default:
					}
					break
				}
			}
		}
	}

	doneTx.Store(true)
	elapsed := time.Since(start).Seconds()
	time.Sleep(drain)
	doneRx.Store(true)
	wg.Wait()

	fmt.Println("---")
	for _, d := range dirs {
		d.sampleDrops()
		s := d.snapshot()
		lost := int64(s.expected) - int64(s.count)
		fmt.Printf("%s total tx %d frames %.2f GB | rx %d frames %.2f GB | lost %d (%.3g%%) crc %d badmagic %d badlen %d kdrop %d\n",
			d.label,
			s.txFrames, float64(s.txBytes)/1e9,
			s.rxFrames, float64(s.rxBytes)/1e9,
			lost, 100*float64(lost)/float64(max(s.expected, 1)),
			s.crcErr, s.badMagic, s.badLen, d.drops)
		fmt.Printf("%s avg tx %.2f Gb/s rx %.2f Gb/s over %.1fs\n",
			d.label,
			gbps(s.txBytes, s.txFrames, elapsed),
			gbps(s.rxBytes, s.rxFrames, elapsed),
			elapsed)
	}
	return nil
}