//go:build linux

package tc

import (
	"fmt"
	"net"

	"github.com/vishvananda/netlink"

	"edge-infra.dev/pkg/lib/kernel/netlink/ip"
)

type MatchIP struct {
	ip   *net.IP
	mask *net.IP
	*ip.IPv4HeaderBitOffset
}

func NewMatchIP() *MatchIP {
	return &MatchIP{}
}

func (m *MatchIP) WithIP(ip *net.IP) *MatchIP {
	m.ip = ip
	return m
}

func (m *MatchIP) WithMask(mask *net.IP) *MatchIP {
	m.mask = mask
	return m
}

func (m *MatchIP) WithIPNet(ipnet *net.IPNet) *MatchIP {
	m.ip = &ipnet.IP
	mask := net.IP(ipnet.Mask)
	m.mask = &mask
	return m
}

func (m *MatchIP) WithIPv4HeaderOffset(offset ip.IPv4HeaderBitOffset) *MatchIP {
	m.IPv4HeaderBitOffset = &offset
	return m
}

type MatchPort struct {
	port uint32
	mask uint32
}

func NewMatchPort() *MatchPort {
	return &MatchPort{}
}

// overrides any existing match port set
func (m *MatchPort) WithDstPort(dstPort uint32) *MatchPort {
	m.port = dstPort
	m.mask = ip.DstPortMask
	return m
}

// overrides any existing match port set
func (m *MatchPort) WithSrcPort(srcPort uint32) *MatchPort {
	m.port = srcPort
	m.mask = ip.SrcPortMask
	return m
}

type U32Filter struct {
	matchIP   []*MatchIP
	matchPort []*MatchPort
	attrs     netlink.FilterAttrs
	classid   uint32
	actions   []netlink.Action
	divisor   uint32
	hashtable uint32
}

func NewU32Filter() *U32Filter {
	return &U32Filter{}
}

func (f *U32Filter) WithHashTable(hashtable uint32) *U32Filter {
	f.hashtable = hashtable
	return f
}

func (f *U32Filter) WithDivisor(divisor uint32) *U32Filter {
	f.divisor = divisor
	return f
}

func (f *U32Filter) WithMatchIP(m *MatchIP) *U32Filter {
	f.matchIP = append(f.matchIP, m)
	return f
}

func (f *U32Filter) WithMatchPort(m *MatchPort) *U32Filter {
	f.matchPort = append(f.matchPort, m)
	return f
}

func (f *U32Filter) WithAttrs(attrs netlink.FilterAttrs) *U32Filter {
	f.attrs = attrs
	return f
}

func (f *U32Filter) WithAction(action netlink.Action) *U32Filter {
	f.actions = append(f.actions, action)
	return f
}

func (f *U32Filter) WithClassID(classid uint32) *U32Filter {
	f.classid = classid
	return f
}

func (f *U32Filter) WithMakeClassID(major, minor uint16) *U32Filter {
	f.classid = netlink.MakeHandle(major, minor)
	return f
}

// atomic filter replace/change
// upstream bug with netlink.FilterReplace https://github.com/vishvananda/netlink/issues/914
func (f *U32Filter) Replace() error {
	filter, err := f.Build()
	if err != nil {
		return err
	}
	link, err := netlink.LinkByIndex(filter.Attrs().LinkIndex)
	if err != nil {
		return err
	}
	if err := netlink.FilterDel(filter); err != nil && !errorIsNotFound(err) {
		return fmt.Errorf(
			"error replacing u32 filter with parent %s and handle %s (dev %s): %w",
			netlink.HandleStr(filter.Attrs().Parent),
			netlink.HandleStr(filter.Attrs().Handle),
			link.Attrs().Name,
			err,
		)
	}
	if err := netlink.FilterAdd(filter); err != nil {
		return fmt.Errorf(
			"error replacing u32 filter with parent %s and handle %s (dev %s): %w",
			netlink.HandleStr(filter.Attrs().Parent),
			netlink.HandleStr(filter.Attrs().Handle),
			link.Attrs().Name,
			err,
		)
	}
	return nil
}

func (f *U32Filter) Add() error {
	filter, err := f.Build()
	if err != nil {
		return err
	}
	if err := netlink.FilterAdd(filter); err != nil && !errorIsFileExists(err) {
		return formatTcError("error creating u32 filter", filter.Attrs().LinkIndex, filter.Attrs().Parent, filter.Attrs().Handle, err)
	}
	return nil
}

func (f *U32Filter) Delete() error {
	filter, err := f.Build()
	if err != nil {
		return err
	}
	if err := netlink.FilterDel(filter); err != nil && !errorIsNotFound(err) {
		return err
	}
	return nil
}

func (f *U32Filter) Build() (netlink.Filter, error) {
	keys := []netlink.TcU32Key{}

	for _, matchIP := range f.matchIP {
		keys = append(keys, getIPSelKey(matchIP))
	}

	for _, matchPort := range f.matchPort {
		keys = append(keys, getPortSelKey(matchPort))
	}

	filter := netlink.U32{}
	sel := netlink.TcU32Sel{
		Flags: netlink.TC_U32_TERMINAL, // avoid printing *flowid https://lore.kernel.org/netdev/20221013084344.0fe8f3de@hermes.local/T/
		Keys:  keys,
	}
	filter.Sel = &sel
	filter.FilterAttrs = f.attrs
	filter.ClassId = f.classid
	filter.Actions = f.actions
	filter.Divisor = f.divisor
	filter.Hash = f.hashtable
	return &filter, nil
}

func getIPSelKey(m *MatchIP) netlink.TcU32Key {
	if m.ip == nil || m.mask == nil || m.IPv4HeaderBitOffset == nil {
		return netlink.TcU32Key{}
	}

	ipAddr := ip.IPv4ToUInt32(*m.ip)
	mask := ip.IPv4ToUInt32(*m.mask)

	offset := int32(*m.IPv4HeaderBitOffset) /* #nosec G115 */

	return netlink.TcU32Key{
		Mask:    mask,
		Val:     ipAddr,
		Off:     offset,
		OffMask: 0,
	}
}

func getPortSelKey(m *MatchPort) netlink.TcU32Key {
	return netlink.TcU32Key{
		Mask:    m.mask,
		Val:     m.port,
		Off:     int32(ip.OptionsBitOffset),
		OffMask: 0,
	}
}