package main

import (
	"encoding/binary"
	"math"
	"runtime"
)

const (
	xtHeader = 32
	ruleHead = 112
)

type tableInfo struct {
	Name      [32]byte
	Valid     uint32
	Hooks     [5]uint32
	Underflow [5]uint32
	Count     uint32
	Size      uint32
}

type replaceHead struct {
	Name      [32]byte
	Valid     uint32
	Count     uint32
	Size      uint32
	Hooks     [5]uint32
	Underflow [5]uint32
	OldCount  uint32
	Address   uint64
}

type ipMatch struct {
	Addresses [4]uint32
	In        [16]byte
	Out       [16]byte
	InMask    [16]byte
	OutMask   [16]byte
	Protocol  uint16
	Flags     uint8
	Inverse   uint8
}

type counter struct {
	Packets uint64
	Bytes   uint64
}

type ruleHeader struct {
	IP       ipMatch
	Cache    uint32
	Target   uint16
	Next     uint16
	From     uint32
	Counters counter
}

type replacement struct {
	data     []byte
	counters []counter
	pins     runtime.Pinner
}

func encode(value any) []byte {
	data, err := binary.Append(nil, binary.NativeEndian, value)
	if err != nil {
		panic(err)
	}
	return data
}

func size16(size int) uint16 {
	if size < 0 || size > math.MaxUint16 {
		panic("iptables size exceeds uint16")
	}
	return uint16(size)
}

func size32(size int) uint32 {
	if size < 0 || uint64(size) > math.MaxUint32 {
		panic("iptables size exceeds uint32")
	}
	return uint32(size)
}

func extension(name string, revision uint8, data []byte) []byte {
	size := (xtHeader + len(data) + 7) &^ 7
	header := struct {
		Size     uint16
		Name     [29]byte
		Revision uint8
	}{Size: size16(size), Revision: revision}
	copy(header.Name[:], name)
	result := append(encode(header), data...)
	return append(result, make([]byte, size-len(result))...)
}

func accept() []byte {
	return extension("", 0, encode(int32(-2)))
}

func redirect(port uint16) []byte {
	rangeInfo := struct {
		Count uint32
		Flags uint32
		IPs   [2]uint32
		Ports [4]byte
	}{Count: 1, Flags: 2}
	binary.BigEndian.PutUint16(rangeInfo.Ports[:2], port)
	binary.BigEndian.PutUint16(rangeInfo.Ports[2:], port)
	return extension("REDIRECT", 0, encode(rangeInfo))
}

func tcp(port uint16) []byte {
	return extension("tcp", 0, encode(struct {
		Source [2]uint16
		Dest   [2]uint16
		Flags  [4]uint8
	}{Source: [2]uint16{0, math.MaxUint16}, Dest: [2]uint16{port, port}}))
}

func owner(uid uint32) []byte {
	return extension("owner", 1, encode(struct {
		UID    [2]uint32
		GID    [2]uint32
		Match  uint8
		Invert uint8
		Pad    [2]byte
	}{UID: [2]uint32{uid, uid}, Match: 1}))
}

func rule(protocol uint16, match, target []byte, loopback bool) []byte {
	header := ruleHeader{
		IP:     ipMatch{Protocol: protocol},
		Target: size16(ruleHead + len(match)),
		Next:   size16(ruleHead + len(match) + len(target)),
	}
	if loopback {
		copy(header.IP.Out[:], "lo")
		header.IP.OutMask[0], header.IP.OutMask[1], header.IP.OutMask[2] = 255, 255, 255
	}
	return append(append(encode(header), match...), target...)
}

func table(value config, count uint32, address uint64) []byte {
	header := replaceHead{Valid: 27, OldCount: count, Address: address}
	copy(header.Name[:], "nat")
	var body []byte
	add := func(data []byte) {
		body = append(body, data...)
		header.Count++
	}
	for _, hook := range []int{0, 1, 3, 4} {
		header.Hooks[hook] = size32(len(body))
		switch hook {
		case 0:
			for _, port := range value.exclude {
				add(rule(6, tcp(port), accept(), false))
			}
			add(rule(6, nil, redirect(value.inbound), false))
		case 3:
			add(rule(0, owner(value.proxyUID), accept(), false))
			add(rule(0, nil, accept(), true))
			add(rule(6, nil, redirect(value.outbound), false))
		}
		header.Underflow[hook] = size32(len(body))
		add(rule(0, nil, accept(), false))
	}
	var label [30]byte
	copy(label[:], "ERROR")
	add(rule(0, nil, extension("ERROR", 0, label[:]), false))
	header.Size = size32(len(body))
	return append(encode(header), body...)
}