package throttler

import (
	"encoding/binary"

	"github.com/vishvananda/netlink"
)

const (
	ipHeaderBytes      = 20
	tcpSrcDstPortBytes = 4
)

func WithDstPort(sel *netlink.TcU32Sel, dport uint16) *netlink.TcU32Sel {
	// ip header(20 bytes)
	// sport(2 bytes), dport(2 bytes)
	buf := make([]byte, tcpSrcDstPortBytes)
	binary.BigEndian.PutUint16(buf[2:], dport)

	key := netlink.TcU32Key{
		Mask: binary.BigEndian.Uint32([]byte{0x00, 0x00, 0xff, 0xff}),
		Val:  binary.BigEndian.Uint32(buf),
		Off:  ipHeaderBytes,
	}

	return addKeyToSel(sel, &key)
}

func WithSrcPort(sel *netlink.TcU32Sel, sport uint16) *netlink.TcU32Sel {
	// ip header(20 bytes)
	// sport(2 bytes), dport(2 bytes)
	buf := make([]byte, tcpSrcDstPortBytes)
	binary.BigEndian.PutUint16(buf[:2], sport)

	key := netlink.TcU32Key{
		Mask: binary.BigEndian.Uint32([]byte{0xff, 0xff, 0x00, 0x00}),
		Val:  binary.BigEndian.Uint32(buf),
		Off:  ipHeaderBytes,
	}

	return addKeyToSel(sel, &key)
}

func addKeyToSel(sel *netlink.TcU32Sel, key *netlink.TcU32Key) *netlink.TcU32Sel {
	if sel == nil {
		sel = &netlink.TcU32Sel{
			Flags: netlink.TC_U32_TERMINAL,
		}
	}

	sel.Keys = append(sel.Keys, *key)
	sel.Nkeys++

	return sel
}