package netutil

import (
	"net"

	"github.com/yl2chen/cidranger"
)

type NetworkTree struct {
	ranger cidranger.Ranger
}

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
}