package main

import (
	"fmt"
	"runtime"
	"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 }

func setBufForce(fd, forceOpt, opt, size int) error {
	if err := unix.SetsockoptInt(fd, unix.SOL_SOCKET, forceOpt, size); err == nil {
		return nil
	}
	return unix.SetsockoptInt(fd, unix.SOL_SOCKET, opt, size)
}

func sockBufSize(fd, opt int) int {
	v, err := unix.GetsockoptInt(fd, unix.SOL_SOCKET, opt)
	if err != nil {
		return -1
	}
	return v
}

func openTxSocket(ifindex, sndbuf 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 := setBufForce(fd, unix.SO_SNDBUFFORCE, unix.SO_SNDBUF, sndbuf); err != nil {
		unix.Close(fd)
		return -1, fmt.Errorf("sndbuf: %w", err)
	}
	return fd, nil
}

func openRxSocket(ifindex, rcvbuf, fanoutID, fanoutMode int) (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 := setBufForce(fd, unix.SO_RCVBUFFORCE, unix.SO_RCVBUF, rcvbuf); 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)
	}
	if fanoutMode >= 0 {
		arg := (fanoutMode << 16) | (fanoutID & 0xffff)
		if err := unix.SetsockoptInt(fd, unix.SOL_PACKET, unix.PACKET_FANOUT, arg); err != nil {
			unix.Close(fd)
			return -1, fmt.Errorf("fanout: %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
}

func pinTo(cpu int) error {
	if cpu < 0 {
		return nil
	}
	runtime.LockOSThread()
	var set unix.CPUSet
	set.Zero()
	set.Set(cpu)
	return unix.SchedSetaffinity(0, &set)
}

func packetDrops(fd int) uint64 {
	st, err := unix.GetsockoptTpacketStats(fd, unix.SOL_PACKET, unix.PACKET_STATISTICS)
	if err != nil {
		return 0
	}
	return uint64(st.Drops)
}