// Package netlink is a partial implementation of the netlink protocol.
package netlink

import (
	"context"
	"fmt"
	"os"
	"syscall"
	"time"
	"unsafe"
)

// net/netlink/af_netlink.c
const maxRecvmsgLen = 32768

const (
	// stateOpen denotes an open conn.
	stateOpen uint32 = 1 << iota
)

// A Conn represents resources associated to a netlink socket.
type Conn struct {
	// AF_NETLINK socket.
	f *os.File
	// For using runtime polling via f.
	raw syscall.RawConn
	// Port ID assigned by the kernel.
	port uint32
	// Internal connection status.
	state uint32
	// Kernel module or netlink group to communicate with.
	family int
	// Message sequence number.
	seq uint32
	// For pending outgoing message.
	typ, flags uint16
	// Outgoing position in buf.
	pos int
	// Pages holding incoming and outgoing messages.
	buf [maxRecvmsgLen]byte
	// An instant some time after conn was established, but before the first
	// I/O operation on f through raw. This serves as a cached deadline to
	// cancel blocking I/O.
	t time.Time
}

// Dial returns the address of a newly connected generic netlink connection of
// specified family and groups.
//
// For a nonzero rcvbuf, the socket receive buffer size is set to its absolute
// value via SO_RCVBUF for a positive value, or SO_RCVBUFFORCE for a negative
// value.
func Dial(family int, groups uint32, rcvbuf int64) (*Conn, error) {
	var c Conn
	if fd, err := syscall.Socket(
		syscall.AF_NETLINK,
		syscall.SOCK_RAW|syscall.SOCK_NONBLOCK|syscall.SOCK_CLOEXEC,
		family,
	); err != nil {
		return nil, os.NewSyscallError("socket", err)
	} else if err = syscall.Bind(fd, &syscall.SockaddrNetlink{
		Family: syscall.AF_NETLINK,
		Groups: groups,
	}); err != nil {
		_ = syscall.Close(fd)
		return nil, os.NewSyscallError("bind", err)
	} else {
		var addr syscall.Sockaddr
		if addr, err = syscall.Getsockname(fd); err != nil {
			_ = syscall.Close(fd)
			return nil, os.NewSyscallError("getsockname", err)
		}
		switch a := addr.(type) {
		case *syscall.SockaddrNetlink:
			c.port = a.Pid

		default: // unreachable
			_ = syscall.Close(fd)
			return nil, syscall.ENOTRECOVERABLE
		}

		if rcvbuf != 0 {
			opt := syscall.SO_RCVBUF
			if rcvbuf < 0 {
				opt = syscall.SO_RCVBUFFORCE
				rcvbuf = -rcvbuf
			}
			if err = syscall.SetsockoptInt(
				fd,
				syscall.SOL_SOCKET,
				opt,
				int(rcvbuf),
			); err != nil {
				_ = syscall.Close(fd)
				return nil, os.NewSyscallError("setsockopt", err)
			}
		}

		c.family = family
		c.f = os.NewFile(uintptr(fd), "netlink")
		if c.raw, err = c.f.SyscallConn(); err != nil {
			_ = c.f.Close()
			return nil, err
		}
		c.state |= stateOpen
	}

	c.pos = syscall.NLMSG_HDRLEN
	c.t = time.Now().UTC()
	return &c, nil
}

// ok returns whether conn is still open.
func (c *Conn) ok() bool { return c.state&stateOpen != 0 }

// Close closes the underlying socket.
func (c *Conn) Close() error {
	if !c.ok() {
		return syscall.EINVAL
	}
	c.state &= ^stateOpen
	return c.f.Close()
}

// Recvmsg wraps recv(2) with nonblocking behaviour via the runtime network poller.
//
// The returned slice is valid until the next call to Recvmsg.
func (c *Conn) Recvmsg(
	ctx context.Context,
	flags int,
) (data []byte, recvflags int, from syscall.Sockaddr, err error) {
	if err = c.f.SetReadDeadline(time.Time{}); err != nil {
		return
	}

	var n int
	data = c.buf[:]

	if ctx == nil {
		rcErr := c.raw.Control(func(fd uintptr) {
			n, _, recvflags, from, err = syscall.Recvmsg(int(fd), data, nil, flags)
		})
		if n >= 0 {
			data = data[:n]
		}

		if err != nil {
			err = os.NewSyscallError("recvmsg", err)
		} else {
			err = rcErr
		}
		return
	}

	done := make(chan error, 1)
	go func() {
		rcErr := c.raw.Read(func(fd uintptr) (done bool) {
			n, _, recvflags, from, err = syscall.Recvmsg(int(fd), data, nil, flags)
			return err != syscall.EWOULDBLOCK
		})
		if n >= 0 {
			data = data[:n]
		}
		done <- rcErr
	}()

	select {
	case rcErr := <-done:
		if err != nil {
			err = os.NewSyscallError("recvmsg", err)
		} else {
			err = rcErr
		}
		return

	case <-ctx.Done():
		cancelErr := c.f.SetReadDeadline(c.t)
		<-done
		if cancelErr != nil {
			err = cancelErr
		} else {
			err = ctx.Err()
		}
		return
	}
}

// Sendmsg wraps send(2) with nonblocking behaviour via the runtime network poller.
func (c *Conn) Sendmsg(
	ctx context.Context,
	p []byte,
	to syscall.Sockaddr,
	flags int,
) (err error) {
	if err = c.f.SetWriteDeadline(time.Time{}); err != nil {
		return
	}

	done := make(chan error, 1)
	go func() {
		done <- c.raw.Write(func(fd uintptr) (done bool) {
			err = syscall.Sendmsg(int(fd), p, nil, to, flags)
			return err != syscall.EWOULDBLOCK
		})
	}()

	select {
	case rcErr := <-done:
		if err != nil {
			err = os.NewSyscallError("sendmsg", err)
		} else {
			err = rcErr
		}
		return

	case <-ctx.Done():
		cancelErr := c.f.SetWriteDeadline(c.t)
		<-done
		if cancelErr != nil {
			err = cancelErr
		} else {
			err = ctx.Err()
		}
		return
	}
}

// Msg is type constraint for types sent over the wire via netlink.
//
// No pointer types or compound types containing pointers may appear here.
type Msg interface {
	syscall.NlMsghdr | syscall.NlMsgerr |
		syscall.IfAddrmsg | RtAttrMsg[InAddr] |
		syscall.IfInfomsg
}

// As returns data as the specified netlink message type.
func As[M Msg](data []byte) *M {
	var v M
	if unsafe.Sizeof(v) != uintptr(len(data)) {
		return nil
	}
	return (*M)(unsafe.Pointer(unsafe.SliceData(data)))
}

// add queues a value to be sent by conn.
func add[M Msg](c *Conn, p *M) bool {
	pos := c.pos
	c.pos += int(unsafe.Sizeof(*p))
	if c.pos > len(c.buf) {
		c.pos = pos
		return false
	}
	*(*M)(unsafe.Pointer(&c.buf[pos])) = *p
	return true
}

// InconsistentError describes a reply from the kernel that is not consistent
// with the internal state tracked by this package.
type InconsistentError struct {
	// Offending header.
	syscall.NlMsghdr
	// Expected message sequence.
	Seq uint32
	// Expected pid.
	Pid uint32
}

func (*InconsistentError) Unwrap() error { return os.ErrInvalid }
func (e *InconsistentError) Error() string {
	s := "netlink socket has inconsistent state"
	switch {
	case e.Seq != e.NlMsghdr.Seq:
		s += fmt.Sprintf(": seq %d != %d", e.Seq, e.NlMsghdr.Seq)
	case e.Pid != e.NlMsghdr.Pid:
		s += fmt.Sprintf(": pid %d != %d", e.Pid, e.NlMsghdr.Pid)
	}
	return s
}

// checkReply checks the message header of a reply from the kernel.
func (c *Conn) checkReply(header *syscall.NlMsghdr) error {
	if header.Seq != c.seq || header.Pid != c.port {
		return &InconsistentError{*header, c.seq, c.port}
	}
	return nil
}

// pending returns the valid slice of buf and initialises pos.
func (c *Conn) pending() []byte {
	buf := c.buf[:c.pos]
	c.pos = syscall.NLMSG_HDRLEN

	*(*syscall.NlMsghdr)(unsafe.Pointer(unsafe.SliceData(buf))) = syscall.NlMsghdr{
		Len:   uint32(len(buf)),
		Type:  c.typ,
		Flags: c.flags,
		Seq:   c.seq,
		Pid:   c.port,
	}
	return buf
}

// Complete indicates the completion of a roundtrip.
type Complete struct{}

// Error returns a hardcoded string that should never be displayed to the user.
func (Complete) Error() string { return "returning from roundtrip" }

// HandlerFunc handles [syscall.NetlinkMessage] and returns a non-nil error to
// discontinue the receiving of more messages.
type HandlerFunc func(resp []syscall.NetlinkMessage) error

// receive receives from a socket with specified flags until a non-nil error is
// returned by f. An error of type [Complete] is returned as nil.
func (c *Conn) receive(ctx context.Context, f HandlerFunc, flags int) error {
	for {
		var resp []syscall.NetlinkMessage
		if data, _, _, err := c.Recvmsg(ctx, flags); err != nil {
			return err
		} else if len(data) < syscall.NLMSG_HDRLEN {
			return syscall.EBADE
		} else if resp, err = syscall.ParseNetlinkMessage(data); err != nil {
			return err
		}

		if err := f(resp); err != nil {
			if err == (Complete{}) {
				return nil
			}
			return err
		}
	}
}

// Roundtrip sends the pending message and handles the reply.
func (c *Conn) Roundtrip(ctx context.Context, f HandlerFunc) error {
	if !c.ok() {
		return syscall.EINVAL
	}
	defer func() { c.seq++ }()

	if err := c.Sendmsg(ctx, c.pending(), &syscall.SockaddrNetlink{
		Family: syscall.AF_NETLINK,
	}, 0); err != nil {
		return err
	}

	return c.receive(ctx, f, 0)
}
