// Copyright 2025 Antonio Ojea
// SPDX-License-Identifier: Apache-2.0

package network

import (
	"fmt"
	"math/bits"
	"net/netip"
)

// iptree implement a radix tree that uses IP prefixes as nodes and allows to store values in each node.
// Example:
//
// 	r := New[int]()
//
//	prefixes := []string{
//		"0.0.0.0/0",
//		"10.0.0.0/8",
//		"10.0.0.0/16",
//		"10.1.0.0/16",
//		"10.1.1.0/24",
//		"10.1.244.0/24",
//		"10.0.0.0/24",
//		"10.0.0.3/32",
//		"192.168.0.0/24",
//		"192.168.0.0/28",
//		"192.168.129.0/28",
//	}
//	for _, k := range prefixes {
//		r.InsertPrefix(netip.MustParsePrefix(k), 0)
//	}
//
// (*) means the node is not public, is not storing any value
//
// 0.0.0.0/0 --- 10.0.0.0/8 --- *10.0.0.0/15 --- 10.0.0.0/16 --- 10.0.0.0/24 --- 10.0.0.3/32
//  |                                 |
//  |                                 \ -------- 10.1.0.0/16 --- 10.1.1.0/24
//  |                                                 |
//  |	                                               \ ------- 10.1.244.0/24
//  |
//  \------ *192.168.0.0/16 --- 192.168.0.0/24 --- 192.168.0.0/28
//                   |
//                    \ -------- 192.168.129.0/28

// node is an element of radix tree with a netip.Prefix optimized to store IP prefixes.
type node[T any] struct {
	// prefix network CIDR
	prefix netip.Prefix
	// public nodes are used to store values
	public bool
	val    T

	child [2]*node[T] // binary tree
}

// mergeChild allow to compress the tree
// when n has exactly one child and no value
// p -> n -> b -> c ==> p -> b -> c
func (n *node[T]) mergeChild() {
	// public nodes can not be merged
	if n.public {
		return
	}
	// can not merge if there are two children
	if n.child[0] != nil &&
		n.child[1] != nil {
		return
	}
	// can not merge if there are no children
	if n.child[0] == nil &&
		n.child[1] == nil {
		return
	}
	// find the child and merge it
	var child *node[T]
	if n.child[0] != nil {
		child = n.child[0]
	} else if n.child[1] != nil {
		child = n.child[1]
	}
	n.prefix = child.prefix
	n.public = child.public
	n.val = child.val
	n.child = child.child
	// remove any references from the deleted node
	// to avoid memory leak
	child.child[0] = nil
	child.child[1] = nil
}

// Tree is a radix tree for IPv4 and IPv6 networks.
type Tree[T any] struct {
	rootV4 *node[T]
	rootV6 *node[T]
}

// NewIPTree creates a new Radix Tree for IP addresses.
func NewIPTree[T any]() *Tree[T] {
	return &Tree[T]{
		rootV4: &node[T]{
			prefix: netip.PrefixFrom(netip.IPv4Unspecified(), 0),
		},
		rootV6: &node[T]{
			prefix: netip.PrefixFrom(netip.IPv6Unspecified(), 0),
		},
	}
}

// GetPrefix returns the stored value and true if the exact prefix exists in the tree.
func (t *Tree[T]) GetPrefix(prefix netip.Prefix) (T, bool) {
	var zeroT T

	n := t.rootV4
	if prefix.Addr().Is6() {
		n = t.rootV6
	}
	bitPosition := 0
	// mask the address for sanity
	address := prefix.Masked().Addr()
	// we can't check longer than the request mask
	mask := prefix.Bits()
	// walk the network bits of the prefix
	for bitPosition < mask {
		// Look for a child checking the bit position after the mask
		n = n.child[getBitFromAddr(address, bitPosition+1)]
		if n == nil {
			return zeroT, false
		}
		// check we are in the right branch comparing the suffixes
		if !n.prefix.Contains(address) {
			return zeroT, false
		}
		// update the new bit position with the new node mask
		bitPosition = n.prefix.Bits()
	}
	// check if this node is a public node and contains a prefix
	if n != nil && n.public && n.prefix == prefix {
		return n.val, true
	}

	return zeroT, false
}

// LongestPrefixMatch returns the longest prefix match, the stored value and true if exist.
// For example, considering the following prefixes 192.168.20.16/28 and 192.168.0.0/16,
// when the address 192.168.20.19/32 is looked up it will return 192.168.20.16/28.
func (t *Tree[T]) LongestPrefixMatch(prefix netip.Prefix) (netip.Prefix, T, bool) {
	n := t.rootV4
	if prefix.Addr().Is6() {
		n = t.rootV6
	}

	var last *node[T]
	// bit position is given by the mask bits
	bitPosition := 0
	// mask the address
	address := prefix.Masked().Addr()
	mask := prefix.Bits()
	// walk the network bits of the prefix
	for bitPosition < mask {
		if n.public {
			last = n
		}
		// Look for a child checking the bit position after the mask
		n = n.child[getBitFromAddr(address, bitPosition+1)]
		if n == nil {
			break
		}
		// check we are in the right branch comparing the suffixes
		if !n.prefix.Contains(address) {
			break
		}
		// update the new bit position with the new node mask
		bitPosition = n.prefix.Bits()
	}

	if n != nil && n.public && n.prefix == prefix {
		last = n
	}

	if last != nil {
		return last.prefix, last.val, true
	}
	var zeroT T
	return netip.Prefix{}, zeroT, false
}

// ShortestPrefixMatch returns the shortest prefix match, the stored value and true if exist.
// For example, considering the following prefixes 192.168.20.16/28 and 192.168.0.0/16,
// when the address 192.168.20.19/32 is looked up it will return 192.168.0.0/16.
func (t *Tree[T]) ShortestPrefixMatch(prefix netip.Prefix) (netip.Prefix, T, bool) {
	var zeroT T

	n := t.rootV4
	if prefix.Addr().Is6() {
		n = t.rootV6
	}
	// bit position is given by the mask bits
	bitPosition := 0
	// mask the address
	address := prefix.Masked().Addr()
	mask := prefix.Bits()
	for bitPosition < mask {
		if n.public {
			return n.prefix, n.val, true
		}
		// Look for a child checking the bit position after the mask
		n = n.child[getBitFromAddr(address, bitPosition+1)]
		if n == nil {
			return netip.Prefix{}, zeroT, false
		}
		// check we are in the right branch comparing the suffixes
		if !n.prefix.Contains(address) {
			return netip.Prefix{}, zeroT, false
		}
		// update the new bit position with the new node mask
		bitPosition = n.prefix.Bits()
	}

	if n != nil && n.public && n.prefix == prefix {
		return n.prefix, n.val, true
	}
	return netip.Prefix{}, zeroT, false
}

// InsertPrefix is used to add a new entry or update
// an existing entry. Returns true if updated.
func (t *Tree[T]) InsertPrefix(prefix netip.Prefix, v T) bool {
	n := t.rootV4
	if prefix.Addr().Is6() {
		n = t.rootV6
	}
	var parent *node[T]
	// bit position is given by the mask bits
	bitPosition := 0
	// mask the address
	address := prefix.Masked().Addr()
	mask := prefix.Bits()
	for bitPosition < mask {
		// Look for a child checking the bit position after the mask
		childIndex := getBitFromAddr(address, bitPosition+1)
		parent = n
		n = n.child[childIndex]
		// if no child create a new one with
		if n == nil {
			parent.child[childIndex] = &node[T]{
				public: true,
				val:    v,
				prefix: prefix,
			}
			return false
		}

		// update the new bit position with the new node mask
		bitPosition = n.prefix.Bits()

		// continue if we are in the right branch and current
		// node is our parent
		if n.prefix.Contains(address) && bitPosition <= mask {
			continue
		}

		// Split the node and add a new child:
		// - Case 1: parent -> child -> n
		// - Case 2: parent -> newnode |--> child
		//                             |--> n
		child := &node[T]{
			prefix: prefix,
			public: true,
			val:    v,
		}
		// Case 1: existing node is a sibling
		if prefix.Contains(n.prefix.Addr()) && bitPosition > mask {
			// parent to child
			parent.child[childIndex] = child
			pos := prefix.Bits() + 1
			// calculate if the sibling is at the left or right
			child.child[getBitFromAddr(n.prefix.Addr(), pos)] = n
			return false
		}

		// Case 2: existing node has the same mask but different base address
		// add common ancestor and branch on it
		ancestor := findAncestor(prefix, n.prefix)
		link := &node[T]{
			prefix: ancestor,
		}
		pos := parent.prefix.Bits() + 1
		parent.child[getBitFromAddr(ancestor.Addr(), pos)] = link
		// ancestor -> children
		pos = ancestor.Bits() + 1
		idxChild := getBitFromAddr(prefix.Addr(), pos)
		idxN := getBitFromAddr(n.prefix.Addr(), pos)
		if idxChild == idxN {
			panic(fmt.Sprintf("wrong ancestor %s: child %s N %s", ancestor.String(), prefix.String(), n.prefix.String()))
		}
		link.child[idxChild] = child
		link.child[idxN] = n
		return false
	}

	// if already exist update it and make it public
	if n != nil && n.prefix == prefix {
		if n.public {
			n.val = v
			n.public = true
			return true
		}
		n.val = v
		n.public = true
		return false
	}

	return false
}

// DeletePrefix delete the exact prefix and return true if it existed.
func (t *Tree[T]) DeletePrefix(prefix netip.Prefix) bool {
	root := t.rootV4
	if prefix.Addr().Is6() {
		root = t.rootV6
	}
	var parent *node[T]
	n := root
	// bit position is given by the mask bits
	bitPosition := 0
	// mask the address
	address := prefix.Masked().Addr()
	mask := prefix.Bits()
	for bitPosition < mask {
		// Look for a child checking the bit position after the mask
		parent = n
		n = n.child[getBitFromAddr(address, bitPosition+1)]
		if n == nil {
			return false
		}
		// check we are in the right branch comparing the suffixes
		if !n.prefix.Contains(address) {
			return false
		}
		// update the new bit position with the new node mask
		bitPosition = n.prefix.Bits()
	}
	// check if the node contains the prefix we want to delete
	if n.prefix != prefix {
		return false
	}
	// Delete the value
	n.public = false
	var zeroT T
	n.val = zeroT

	nodeChildren := 0
	if n.child[0] != nil {
		nodeChildren++
	}
	if n.child[1] != nil {
		nodeChildren++
	}
	// If there is a parent and this node does not have any children
	// this is a leaf so we can delete this node.
	// - parent -> child(to be deleted)
	if parent != nil && nodeChildren == 0 {
		// nolint:gocritic
		if parent.child[0] != nil && parent.child[0] == n {
			parent.child[0] = nil
		} else if parent.child[1] != nil && parent.child[1] == n {
			parent.child[1] = nil
		} else {
			panic("wrong parent")
		}
		n = nil
	}
	// Check if we should merge this node
	// The root node can not be merged
	if n != root && nodeChildren == 1 {
		n.mergeChild()
	}
	// Check if we should merge the parent's other child
	// parent -> deletedNode
	//        |--> child
	parentChildren := 0
	if parent != nil {
		if parent.child[0] != nil {
			parentChildren++
		}
		if parent.child[1] != nil {
			parentChildren++
		}
		if parent != root && parentChildren == 1 && !parent.public {
			parent.mergeChild()
		}
	}
	return true
}

// for testing, returns the number of public nodes in the tree.
func (t *Tree[T]) Len(isV6 bool) int {
	count := 0
	t.DepthFirstWalk(isV6, func(k netip.Prefix, v T) bool {
		count++
		return false
	})
	return count
}

// WalkFn is used when walking the tree. Takes a
// key and value, returning if iteration should
// be terminated.
type WalkFn[T any] func(s netip.Prefix, v T) bool

// DepthFirstWalk is used to walk the tree of the corresponding IP family
func (t *Tree[T]) DepthFirstWalk(isIPv6 bool, fn WalkFn[T]) {
	if isIPv6 {
		recursiveWalk(t.rootV6, fn)
	}
	recursiveWalk(t.rootV4, fn)
}

// recursiveWalk is used to do a pre-order walk of a node
// recursively. Returns true if the walk should be aborted
func recursiveWalk[T any](n *node[T], fn WalkFn[T]) bool {
	if n == nil {
		return true
	}
	// Visit the public values if any
	if n.public && fn(n.prefix, n.val) {
		return true
	}

	// Recurse on the children
	if n.child[0] != nil {
		if recursiveWalk(n.child[0], fn) {
			return true
		}
	}
	if n.child[1] != nil {
		if recursiveWalk(n.child[1], fn) {
			return true
		}
	}
	return false
}

// WalkPrefix is used to walk the tree under a prefix
func (t *Tree[T]) WalkPrefix(prefix netip.Prefix, fn WalkFn[T]) {
	n := t.rootV4
	if prefix.Addr().Is6() {
		n = t.rootV6
	}
	bitPosition := 0
	// mask the address for sanity
	address := prefix.Masked().Addr()
	// we can't check longer than the request mask
	mask := prefix.Bits()
	// walk the network bits of the prefix
	for bitPosition < mask {
		// Look for a child checking the bit position after the mask
		n = n.child[getBitFromAddr(address, bitPosition+1)]
		if n == nil {
			return
		}
		// check we are in the right branch comparing the suffixes
		if !n.prefix.Contains(address) {
			break
		}
		// update the new bit position with the new node mask
		bitPosition = n.prefix.Bits()
	}
	recursiveWalk[T](n, fn)

}

// WalkPath is used to walk the tree, but only visiting nodes
// from the root down to a given IP prefix. Where WalkPrefix walks
// all the entries *under* the given prefix, this walks the
// entries *above* the given prefix.
func (t *Tree[T]) WalkPath(path netip.Prefix, fn WalkFn[T]) {
	n := t.rootV4
	if path.Addr().Is6() {
		n = t.rootV6
	}
	bitPosition := 0
	// mask the address for sanity
	address := path.Masked().Addr()
	// we can't check longer than the request mask
	mask := path.Bits()
	// walk the network bits of the prefix
	for bitPosition < mask {
		// Visit the public values if any
		if n.public && fn(n.prefix, n.val) {
			return
		}
		// Look for a child checking the bit position after the mask
		n = n.child[getBitFromAddr(address, bitPosition+1)]
		if n == nil {
			return
		}
		// check we are in the right branch comparing the suffixes
		if !n.prefix.Contains(address) {
			return
		}
		// update the new bit position with the new node mask
		bitPosition = n.prefix.Bits()
	}
	// check if this node is a public node and contains a prefix
	if n != nil && n.public && n.prefix == path {
		fn(n.prefix, n.val)
	}
}

// TopLevelPrefixes is used to return a map with all the Top Level prefixes
// from the corresponding IP family and its values.
// For example, if the tree contains entries for 10.0.0.0/8, 10.1.0.0/16, and 192.168.0.0/16,
// this will return 10.0.0.0/8 and 192.168.0.0/16.
func (t *Tree[T]) TopLevelPrefixes(isIPv6 bool) map[netip.Prefix]T {
	if isIPv6 {
		return t.topLevelPrefixes(t.rootV6)
	}
	return t.topLevelPrefixes(t.rootV4)
}

// topLevelPrefixes is used to return a map with all the Top Level prefixes and its values
func (t *Tree[T]) topLevelPrefixes(root *node[T]) map[netip.Prefix]T {
	result := map[netip.Prefix]T{}
	queue := []*node[T]{root}

	for len(queue) > 0 {
		n := queue[0]
		queue = queue[1:]
		// store and continue, only interested on the top level prefixes
		if n.public {
			result[n.prefix] = n.val
			continue
		}
		if n.child[0] != nil {
			queue = append(queue, n.child[0])
		}
		if n.child[1] != nil {
			queue = append(queue, n.child[1])
		}
	}
	return result
}

// GetHostIPPrefixMatches returns the list of prefixes that contain the specified Host IP.
// An IP is considered a Host IP if is within the subnet range and is not the network address
// or, if IPv4, the broadcast address (RFC 1878).
func (t *Tree[T]) GetHostIPPrefixMatches(ip netip.Addr) map[netip.Prefix]T {
	// walk the tree to find all the prefixes containing this IP
	ipPrefix := netip.PrefixFrom(ip, ip.BitLen())
	prefixes := map[netip.Prefix]T{}
	t.WalkPath(ipPrefix, func(k netip.Prefix, v T) bool {
		if prefixContainIP(k, ipPrefix.Addr()) {
			prefixes[k] = v
		}
		return false
	})
	return prefixes
}

// assume starts at 0 from the MSB: 0.1.2......31
// return 0 or 1
func getBitFromAddr(ip netip.Addr, pos int) int {
	bytes := ip.AsSlice()
	// get the byte in the slice
	index := (pos - 1) / 8
	if index >= len(bytes) {
		panic(fmt.Sprintf("ip %s pos %d index %d bytes %v", ip, pos, index, bytes))
	}
	// get the offset inside the byte
	offset := (pos - 1) % 8
	// check if the bit is set
	if bytes[index]&(uint8(0x80)>>offset) > 0 {
		return 1
	}
	return 0
}

// find the common subnet, aka the one with the common prefix
func findAncestor(a, b netip.Prefix) netip.Prefix {
	bytesA := a.Addr().AsSlice()
	bytesB := b.Addr().AsSlice()
	bytes := make([]byte, len(bytesA))

	max := a.Bits()
	if l := b.Bits(); l < max {
		max = l
	}

	mask := 0
	for i := range bytesA {
		xor := bytesA[i] ^ bytesB[i]
		if xor == 0 {
			bytes[i] = bytesA[i]
			mask += 8

		} else {
			pos := bits.LeadingZeros8(xor)
			mask += pos
			// mask off the non leading zeros
			bytes[i] = bytesA[i] & (^uint8(0) << (8 - pos))
			break
		}
	}
	if mask > max {
		mask = max
	}

	addr, ok := netip.AddrFromSlice(bytes)
	if !ok {
		panic(bytes)
	}
	ancestor := netip.PrefixFrom(addr, mask)
	return ancestor.Masked()
}

// prefixContainIP returns true if the given IP is contained with the prefix,
// is not the network address and also, if IPv4, is not the broadcast address.
// This is required because the Kubernetes allocators reserve these addresses
// so IPAddresses can not block deletion of this ranges.
func prefixContainIP(prefix netip.Prefix, ip netip.Addr) bool {
	// if the IP is the network address is not contained
	if prefix.Masked().Addr() == ip {
		return false
	}
	// the broadcast address is not considered contained for IPv4
	if !ip.Is6() {
		ipLast, err := broadcastAddress(prefix)
		if err != nil || ipLast == ip {
			return false
		}
	}
	return prefix.Contains(ip)
}
