package wg

import (
	"encoding/binary"
	"errors"
	"net/netip"

	"github.com/hashicorp/go-set/v3"
)

type IPPool struct {
	scope   netip.Prefix
	leasure *set.Set[netip.Addr]
	cache   netip.Addr
}

var ErrIPPoolExhausted = errors.New("IP Pool Exhausted")

// @param base_ip4: is used as the gateway address
func IPPool_Init(base_ip4 netip.Addr, cidr uint8) *IPPool {
	scope := netip.PrefixFrom(base_ip4, int(cidr))
	pool := &IPPool{leasure: set.New[netip.Addr](1 << (32 - cidr)), scope: scope}
	pool.cache = pool.scope.Addr().Next()

	//calculate broadcast address
	var wildcard uint32 = (1 << (32 - cidr)) - 1

	wildcard_mask := [4]byte{}
	binary.BigEndian.PutUint32(wildcard_mask[:4], wildcard)
	broadcast := base_ip4.As4()
	broadcast[0] |= wildcard_mask[0]
	broadcast[1] |= wildcard_mask[1]
	broadcast[2] |= wildcard_mask[2]
	broadcast[3] |= wildcard_mask[3]

	broadcast_addr := netip.AddrFrom4(broadcast)
	pool.Lease(broadcast_addr)
	//gateway
	pool.Lease(base_ip4)

	return pool
}
func (p *IPPool) Next() (netip.Addr, error) {
	cyclic := false
	for ; ; p.cache = p.cache.Next() {
		if p.leasure.Contains(p.cache) && p.cache.IsValid() {
			continue
		} else {
			if p.cache.IsValid() && p.scope.Contains(p.cache) && !p.leasure.Contains(p.cache) {
				return p.cache, nil
			} else if !cyclic {
				p.cache = p.scope.Masked().Addr()
				cyclic = true

			} else {
				return p.cache, ErrIPPoolExhausted
			}
		}
	}
}

func (p *IPPool) Lease(addr netip.Addr) {
	p.leasure.Insert(addr)
}
func (p *IPPool) HasIP(addr netip.Addr) {
	p.leasure.Contains(addr)
}