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
	tag   string
	idx   int
	mac   [6]byte
	mtu   int
	speed float64
}

func (e endpoint) macString() string {
	return fmt.Sprintf("%02x:%02x:%02x:%02x:%02x:%02x",
		e.mac[0], e.mac[1], e.mac[2], e.mac[3], e.mac[4], e.mac[5])
}

type direction struct {
	label   string
	short   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
	nicTX0 nicCounters
	nicRX0 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)
	e := endpoint{name: name, idx: ifi.Index, mac: mac, mtu: ifi.MTU}
	if v, ok := readUint("/sys/class/net/" + name + "/speed"); ok {
		e.speed = float64(v) / 1000
	}
	return e, 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 {
		expected := d.streams[i].maxSeq.Load() + 1
		c := d.streams[i].count.Load()
		if c == 0 {
			continue
		}
		s.count += c
		s.expected += expected
	}
	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
}

var intervalCols = []colSpec{
	{title: "TIME", width: 6, right: true},
	{title: "DIR", width: 5},
	{title: "TX pps", width: 9, right: true},
	{title: "TX Gb/s", width: 7, right: true},
	{title: "RX pps", width: 9, right: true},
	{title: "RX Gb/s", width: 7, right: true},
	{title: "GAP", width: 9, right: true},
	{title: "CRC", width: 5, right: true},
	{title: "BADMAG", width: 6, right: true},
	{title: "KDROP", width: 7, right: true},
}

func (d *direction) intervalRow(elapsed, secs, target 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

	before := d.drops
	d.sampleDrops()

	return []string{
		fmt.Sprintf("%.0fs", elapsed),
		paint(d.short, cCyan),
		commas(uint64(float64(txF) / secs)),
		rateCell(gbps(txB, txF, secs), target),
		commas(uint64(float64(rxF) / secs)),
		rateCell(gbps(rxB, rxF, secs), target),
		gapCell(int64(now.expected) - int64(now.count)),
		statusCell(now.crcErr - p.crcErr),
		statusCell(now.badMagic - p.badMagic),
		statusCell(d.drops - before),
	}
}

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, paint("  ▲ "+d.tx.name+" tx-side: "+s, cYellow))
	}
	if s := rx.diff(d.nicRX); s != "" {
		parts = append(parts, paint("  ▲ "+d.rx.name+" rx-side: "+s, cYellow))
	}
	d.nicTX = tx
	d.nicRX = rx
	return parts
}

func (d *direction) nicRunTotals() string {
	var parts []string
	if s := readNIC(d.tx.name).diff(d.nicTX0); s != "" {
		parts = append(parts, d.tx.name+" tx-side: "+s)
	}
	if s := readNIC(d.rx.name).diff(d.nicRX0); s != "" {
		parts = append(parts, d.rx.name+" rx-side: "+s)
	}
	return strings.Join(parts, "; ")
}

func buildDirection(label string, tx, rx endpoint, patIdx int, sizes []int, cfg config) (*direction, error) {
	d := &direction{
		label:   label,
		short:   tx.tag + "→" + rx.tag,
		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)
	d.nicTX0 = d.nicTX
	d.nicRX0 = d.nicRX

	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)
		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.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),
			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
	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")
		interval = flag.Duration("interval", time.Second, "report interval")
		fanout   = flag.String("fanout", "lb", "rx fanout mode: none, hash, lb, cpu, rollover")
		duplex   = flag.Bool("duplex", true, "run both directions simultaneously")
		txCPUs   = flag.String("txcpus", "", "comma-separated CPUs to pin tx workers to")
		rxCPUs   = flag.String("rxcpus", "", "comma-separated CPUs to pin rx workers to")
	)
	flag.Parse()

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

const drainTime = 500 * time.Millisecond

func run(aName, bName, sizesArg, patArg, fanout, txCPUsArg, rxCPUsArg string,
	duration, interval time.Duration, txN, rxN, batch int, duplex bool) 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)
			}
		}
	}

	a.tag, b.tag = "A", "B"
	ifnames := []string{a.name, b.name}

	var fatal []string
	var tuneRows [][]string
	for _, r := range configureSystem(ifnames) {
		tuneRows = append(tuneRows, []string{r.item, r.status(), r.detail()})
		if r.fatal {
			fatal = append(fatal, r.item)
		}
	}
	fmt.Println(renderBox("HOST SETTINGS",
		[]string{"CHECK", "STATUS", "DETAIL"},
		[]bool{false, false, false}, tuneRows))
	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,
		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()
		}
	}()

	var linkRows [][]string
	for _, e := range []endpoint{a, b} {
		linkRows = append(linkRows, []string{
			paint(e.tag, cCyan), e.name, e.macString(),
			fmt.Sprintf("%.0f Gb/s", e.speed), fmt.Sprintf("%d", e.mtu),
		})
	}
	fmt.Println(renderBox("LINKS",
		[]string{"TAG", "INTERFACE", "MAC", "SPEED", "MTU"},
		[]bool{false, false, false, true, true}, linkRows))

	target := a.speed
	if target <= 0 {
		target = 10
	}
	sizeStrs := make([]string, len(sizes))
	for i, s := range sizes {
		sizeStrs[i] = fmt.Sprintf("%d", s)
	}
	fmt.Println(renderBox("TEST",
		[]string{"SETTING", "VALUE"},
		[]bool{false, false}, [][]string{
			{"pattern", patterns[patIdx].name},
			{"frame sizes", strings.Join(sizeStrs, " ")},
			{"ethertype", fmt.Sprintf("0x%04x", etherType)},
			{"workers", fmt.Sprintf("%d tx, %d rx per direction", txN, rxN)},
			{"batch", fmt.Sprintf("%d frames per syscall", batch)},
			{"payload verify", "crc32c on every frame"},
			{"rx fanout", fanout},
			{"duplex", fmt.Sprintf("%v", duplex)},
			{"socket buffers", fmt.Sprintf("sndbuf %s, rcvbuf %s (granted)",
				humanBytes(uint64(sockBufSize(dirs[0].txFDs[0], unix.SO_SNDBUF))),
				humanBytes(uint64(sockBufSize(dirs[0].rxFDs[0], unix.SO_RCVBUF))))},
		}))
	fmt.Println(paint("GAP counts frames below the high-water mark not yet drained from the rx queues;", cDim))
	fmt.Println(paint("it is provisional and settles at the drain. RESULTS at the end is authoritative.", cDim))
	fmt.Println()

	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()
	stats := &streamTable{cols: intervalCols, headerEvery: 20}
loop:
	for {
		select {
		case <-sig:
			break loop
		case <-deadline:
			break loop
		case now := <-tick.C:
			secs := now.Sub(last).Seconds()
			last = now
			elapsed := now.Sub(start).Seconds()
			for _, d := range dirs {
				for _, line := range stats.emit(d.intervalRow(elapsed, secs, target)) {
					fmt.Println(line)
				}
				for _, line := range d.reportNIC() {
					fmt.Println(line)
				}
				for {
					select {
					case msg := <-d.reports:
						fmt.Println(paint("  ✗ "+d.short+": "+msg, cRed))
						continue
					default:
					}
					break
				}
			}
		}
	}

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

	fmt.Println()
	var resultRows, nicRows [][]string
	var problems []string
	for _, d := range dirs {
		d.sampleDrops()
		s := d.snapshot()
		lost := int64(s.expected) - int64(s.count)
		lostPct := 100 * float64(lost) / float64(max(s.expected, 1))

		resultRows = append(resultRows, []string{
			paint(d.short, cCyan),
			commas(s.txFrames),
			commas(s.rxFrames),
			humanBytes(s.rxBytes),
			rateCell(gbps(s.txBytes, s.txFrames, elapsed), target),
			rateCell(gbps(s.rxBytes, s.rxFrames, elapsed), target),
			lostCell(lost),
			fmt.Sprintf("%.4f", lostPct),
			statusCell(s.crcErr),
			statusCell(s.badMagic),
			statusCell(s.badLen),
			statusCell(d.drops),
		})

		if n := d.nicRunTotals(); n != "" {
			nicRows = append(nicRows, []string{paint(d.short, cCyan), n})
		}
		if lost != 0 {
			problems = append(problems, fmt.Sprintf("%s lost %d frames (%.4f%%)", d.short, lost, lostPct))
		}
		if s.crcErr != 0 {
			problems = append(problems, fmt.Sprintf("%s had %d payload CRC failures", d.short, s.crcErr))
		}
		if d.drops != 0 {
			problems = append(problems, fmt.Sprintf("%s dropped %d frames in the kernel (host too slow, not the cable)", d.short, d.drops))
		}
	}

	fmt.Println(renderBox(fmt.Sprintf("RESULTS after %.1fs", elapsed),
		[]string{"DIR", "TX FRAMES", "RX FRAMES", "RX DATA", "TX Gb/s", "RX Gb/s", "LOST", "LOST %", "CRC", "BAD", "BADLEN", "KDROP"},
		[]bool{false, true, true, true, true, true, true, true, true, true, true, true},
		resultRows))

	if len(nicRows) > 0 {
		fmt.Println(renderBox("NIC COUNTER CHANGES DURING RUN",
			[]string{"DIR", "COUNTERS"}, []bool{false, false}, nicRows))
		problems = append(problems, "NIC error counters moved during the run")
	}

	fmt.Println()
	if len(problems) == 0 {
		fmt.Println(paint(" PASS ", cBold+"\x1b[42m\x1b[30m") + " " +
			paint(fmt.Sprintf("every frame arrived, no errors, %s each way at %.2f Gb/s",
				humanBytes(dirs[0].snapshot().rxBytes), target), cGreen))
	} else {
		fmt.Println(paint(" FAIL ", cBold+"\x1b[41m\x1b[37m"))
		for _, p := range problems {
			fmt.Println(paint("  ✗ "+p, cRed))
		}
	}
	return nil
}