package main

import (
	"fmt"
	"unsafe"

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

type mmsghdr struct {
	hdr unix.Msghdr
	len uint32
	_   [4]byte
}

func htons(v uint16) uint16 { return v<<8 | v>>8 }

// Set through the FORCE options alone: the plain ones are clamped to wmem_max
// and rmem_max, so falling back to them would quietly leave a fraction of this
// and go on measuring as though it had not.
const (
	sndbufBytes = 8 << 20
	rcvbufBytes = 64 << 20
)

func openTxSocket(ifindex int) (int, error) {
	fd, err := unix.Socket(unix.AF_PACKET, unix.SOCK_RAW, 0)
	if err != nil {
		return -1, fmt.Errorf("socket: %w", err)
	}
	if err := unix.Bind(fd, &unix.SockaddrLinklayer{Ifindex: ifindex}); err != nil {
		unix.Close(fd)
		return -1, fmt.Errorf("bind ifindex %d: %w", ifindex, err)
	}
	if err := unix.SetsockoptInt(fd, unix.SOL_PACKET, unix.PACKET_QDISC_BYPASS, 1); err != nil {
		unix.Close(fd)
		return -1, fmt.Errorf("qdisc bypass: %w", err)
	}
	if err := unix.SetsockoptInt(fd, unix.SOL_SOCKET, unix.SO_SNDBUFFORCE, sndbufBytes); err != nil {
		unix.Close(fd)
		return -1, fmt.Errorf("sndbuf: %w", err)
	}
	return fd, nil
}

func openRxSocket(ifindex int, etherType uint16) (int, error) {
	proto := int(htons(etherType))
	fd, err := unix.Socket(unix.AF_PACKET, unix.SOCK_RAW, proto)
	if err != nil {
		return -1, fmt.Errorf("socket: %w", err)
	}
	if err := unix.SetsockoptInt(fd, unix.SOL_SOCKET, unix.SO_RCVBUFFORCE, rcvbufBytes); err != nil {
		unix.Close(fd)
		return -1, fmt.Errorf("rcvbuf: %w", err)
	}
	if err := unix.Bind(fd, &unix.SockaddrLinklayer{
		Protocol: htons(etherType),
		Ifindex:  ifindex,
	}); err != nil {
		unix.Close(fd)
		return -1, fmt.Errorf("bind ifindex %d: %w", ifindex, err)
	}
	if err := unix.SetsockoptInt(fd, unix.SOL_PACKET, unix.PACKET_IGNORE_OUTGOING, 1); err != nil {
		unix.Close(fd)
		return -1, fmt.Errorf("ignore outgoing: %w", err)
	}
	tv := unix.Timeval{Sec: 0, Usec: 200000}
	if err := unix.SetsockoptTimeval(fd, unix.SOL_SOCKET, unix.SO_RCVTIMEO, &tv); err != nil {
		unix.Close(fd)
		return -1, fmt.Errorf("rcvtimeo: %w", err)
	}
	return fd, nil
}

func sendmmsg(fd int, hdrs []mmsghdr) (int, error) {
	n, _, errno := unix.Syscall6(unix.SYS_SENDMMSG, uintptr(fd),
		uintptr(unsafe.Pointer(&hdrs[0])), uintptr(len(hdrs)), 0, 0, 0)
	if errno != 0 {
		return int(n), errno
	}
	return int(n), nil
}

func recvmmsg(fd int, hdrs []mmsghdr, flags int) (int, error) {
	n, _, errno := unix.Syscall6(unix.SYS_RECVMMSG, uintptr(fd),
		uintptr(unsafe.Pointer(&hdrs[0])), uintptr(len(hdrs)), uintptr(flags), 0, 0)
	if errno != 0 {
		return int(n), errno
	}
	return int(n), nil
}

func newMmsghdrs(bufs [][]byte) ([]mmsghdr, []unix.Iovec) {
	hdrs := make([]mmsghdr, len(bufs))
	iovs := make([]unix.Iovec, len(bufs))
	for i := range bufs {
		iovs[i].Base = &bufs[i][0]
		iovs[i].Len = uint64(len(bufs[i]))
		hdrs[i].hdr.Iov = &iovs[i]
		hdrs[i].hdr.Iovlen = 1
	}
	return hdrs, iovs
}

// Room for one SCM_TIMESTAMPING and its three timespecs.
const cmsgLen = 128

// Receive headers carry a control buffer each, so the mac's receive stamp comes
// back alongside every frame.
func newRxMmsghdrs(bufs [][]byte) ([]mmsghdr, [][]byte) {
	hdrs, _ := newMmsghdrs(bufs)
	oob := make([][]byte, len(bufs))
	for i := range bufs {
		oob[i] = make([]byte, cmsgLen)
		hdrs[i].hdr.Control = &oob[i][0]
		hdrs[i].hdr.Controllen = cmsgLen
	}
	return hdrs, oob
}

func packetDrops(fd int) uint64 {
	st, err := unix.GetsockoptTpacketStats(fd, unix.SOL_PACKET, unix.PACKET_STATISTICS)
	if err != nil {
		panic(fmt.Sprintf("reading packet drop statistics: %v", err))
	}
	return uint64(st.Drops)
}