package main

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

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

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
	specs   []*frameSpec
	txStats []*txStats
	rxStats []*rxStats
	streams []lossWindow
	txFDs   []int
	rxFDs   []int
	reports chan string

	prev     sample
	drops    uint64
	errBase  sample
	dropBase uint64
	nicTX    nicCounters
	nicRX    nicCounters
}

type sample struct {
	txFrames, txBytes uint64
	rxFrames, rxBytes uint64
	lost, late        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 (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 {
		s.lost += d.streams[i].lost.Load()
		s.late += d.streams[i].late.Load()
	}
	return s
}

// Counters keep climbing in the workers, so resetting just moves the origin
// the display subtracts from.
func (d *direction) resetErrors() {
	d.sampleDrops()
	d.errBase = d.snapshot()
	d.dropBase = d.drops
}

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: "UPTIME", width: 9, 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: "LOST", width: 11, right: true},
	{title: "LATE", width: 9, right: true},
	{title: "CRC", width: 7, right: true},
	{title: "BADMAG", width: 7, right: true},
	{title: "KDROP", width: 11, right: true},
	{title: "ERRORS", width: 11, right: true},
}

// One interval's numbers, shared by the console table and the framebuffer so
// both always show the same figures.
type view struct {
	txPPS, rxPPS     float64
	txGbps, rxGbps   float64
	txFrames, txSent uint64
	rxFrames, rxGot  uint64
	lost, late       uint64
	crc, badMagic    uint64
	kdrop, errors    uint64
}

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

	txF := now.txFrames - p.txFrames
	rxF := now.rxFrames - p.rxFrames
	b := d.errBase
	v := view{
		txPPS:    float64(txF) / secs,
		rxPPS:    float64(rxF) / secs,
		txGbps:   gbps(now.txBytes-p.txBytes, txF, secs),
		rxGbps:   gbps(now.rxBytes-p.rxBytes, rxF, secs),
		txFrames: now.txFrames,
		txSent:   now.txBytes,
		rxFrames: now.rxFrames,
		rxGot:    now.rxBytes,
		lost:     now.lost - b.lost,
		late:     now.late - b.late,
		crc:      now.crcErr - b.crcErr,
		badMagic: now.badMagic - b.badMagic,
		kdrop:    d.drops - d.dropBase,
	}
	v.errors = v.lost + v.crc + v.badMagic + (now.badLen - b.badLen) + v.kdrop
	return v
}

func (d *direction) row(elapsed time.Duration, v view, target float64) []string {
	return []string{
		uptime(elapsed),
		paint(d.short, cCyan),
		commas(uint64(v.txPPS)),
		rateCell(v.txGbps, target),
		commas(uint64(v.rxPPS)),
		rateCell(v.rxGbps, target),
		statusCell(v.lost),
		statusCell(v.late),
		statusCell(v.crc),
		statusCell(v.badMagic),
		statusCell(v.kdrop),
		statusCell(v.errors),
	}
}

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 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,
		streams: newLossWindows(cfg.streams),
		reports: make(chan string, 64),
	}
	d.nicTX = readNIC(tx.name)
	d.nicRX = readNIC(rx.name)

	for i := 0; i < cfg.streams; i++ {
		et := uint16(etherBase + i)
		d.specs = append(d.specs, newFrameSpec(patIdx, rx.mac, tx.mac, et, sizes))

		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{})

		fd, err = openRxSocket(rx.idx, et)
		if err != nil {
			return nil, fmt.Errorf("%s rx socket for 0x%04x: %w", label, et, 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.specs[i],
			batch:   cfg.batch,
			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,
			spec:    d.specs[i],
			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 {
	streams int
	batch   int
}

func main() {
	var (
		aName    = flag.String("a", "", "first interface")
		bName    = flag.String("b", "", "second interface")
		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")
		streams  = flag.Int("streams", 7, "independent streams per direction, capped by rx rings; each gets its own ethertype, steered by a flow rule to its own rx queue")
		batch    = flag.Int("batch", 64, "frames per sendmmsg/recvmmsg call")
		duplex   = flag.Bool("duplex", true, "run both directions simultaneously")
	)
	flag.Parse()

	if err := run(*aName, *bName, *sizesArg, *patArg,
		*streams, *batch, *duplex); err != nil {
		fmt.Fprintln(os.Stderr, "error:", err)
		os.Exit(1)
	}
}

const reportInterval = 500 * time.Millisecond

func run(aName, bName, sizesArg, patArg string,
	nStreams, 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
	}
	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}

	if nStreams < 1 {
		return fmt.Errorf("need at least one stream")
	}
	ethertypes := make([]uint16, nStreams)
	for i := range ethertypes {
		ethertypes[i] = uint16(etherBase + i)
	}

	var fatal []string
	var tuneRows [][]string
	for _, r := range configureSystem(ifnames, ethertypes) {
		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{
		streams: nStreams,
		batch:   batch,
	}

	var dirs []*direction
	d0, err := buildDirection(a.name+"->"+b.name, a, b, patIdx, sizes, cfg)
	if err != nil {
		return err
	}
	dirs = append(dirs, d0)
	if duplex {
		d1, err := buildDirection(b.name+"->"+a.name, b, a, patIdx, sizes, cfg)
		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, " ")},
			{"streams", fmt.Sprintf("%d per direction, ethertypes 0x%04x-0x%04x, one rx queue each",
				nStreams, ethertypes[0], ethertypes[len(ethertypes)-1])},
			{"batch", fmt.Sprintf("%d frames per syscall", batch)},
			{"payload verify", "crc32c on every frame"},
			{"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("rates are per interval; error counts are cumulative, press space to reset them", 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)

	space, restoreTerm := watchSpace()
	defer restoreTerm()

	disp, err := newDisplay()
	if err != nil {
		return fmt.Errorf("display: %w", err)
	}
	defer disp.close()

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

	last := time.Now()
	stats := &streamTable{cols: intervalCols, headerEvery: 20}
	for {
		select {
		case <-sig:
			doneTx.Store(true)
			doneRx.Store(true)
			wg.Wait()
			return nil
		case <-space:
			for _, d := range dirs {
				d.resetErrors()
			}
			stats.sinceHeader = 0
			fmt.Println(stats.rule("error counts reset"))
		case now := <-tick.C:
			secs := now.Sub(last).Seconds()
			last = now
			elapsed := now.Sub(start)
			views := make([]view, len(dirs))
			for i, d := range dirs {
				views[i] = d.view(secs)
				for _, line := range stats.emit(d.row(elapsed, views[i], 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
				}
			}
			disp.render(dirs, views, elapsed, target)
		}
	}
}