package dataset

import (
	"net"
	"net/netip"

	"git.maze.io/maze/styx/dataset/nettrie"
	"github.com/yl2chen/cidranger"
)

type NetworkTrie struct {
	*nettrie.Trie
}

func NewNetworkTrie(prefixes ...netip.Prefix) *NetworkTrie {
	trie := &NetworkTrie{
		Trie: nettrie.New(),
	}
	for _, prefix := range prefixes {
		trie.Insert(prefix)
	}
	return trie
}

func (trie *NetworkTrie) ContainsIP(ip net.IP) bool {
	if ip == nil {
		return false
	}
	addr, ok := netip.AddrFromSlice(ip)
	if !ok {
		return false
	}
	return trie.Contains(addr)
}

type NetworkTree struct {
	ranger cidranger.Ranger
}

func MustNetworkTree(networks ...string) *NetworkTree {
	tree, err := NewNetworkTree(networks...)
	if err != nil {
		panic(err)
	}
	return tree
}

func NewNetworkTree(networks ...string) (*NetworkTree, error) {
	tree := &NetworkTree{
		ranger: cidranger.NewPCTrieRanger(),
	}
	for _, cidr := range networks {
		if err := tree.AddCIDR(cidr); err != nil {
			return nil, err
		}
	}
	return tree, nil
}

func (tree *NetworkTree) Add(ipnet *net.IPNet) {
	if ipnet == nil {
		return
	}
	tree.ranger.Insert(cidranger.NewBasicRangerEntry(*ipnet))
}

func (tree *NetworkTree) AddCIDR(cidr string) error {
	_, ipnet, err := net.ParseCIDR(cidr)
	if err != nil {
		return err
	}
	tree.ranger.Insert(cidranger.NewBasicRangerEntry(*ipnet))
	return nil
}

func (tree *NetworkTree) Contains(ip net.IP) bool {
	contains, _ := tree.ranger.Contains(ip)
	return contains
}