package netalloc

import (
	"errors"
	"net"

	"github.com/mikioh/ipaddr"
)

var (
	ErrPrefixFull = errors.New("no more ip addresses available in prefix")
)

type Allocator interface {
	Alloc() (*net.IPAddr, error)
	AllocCIDR() (cidr string, err error)
	Free(*net.IPAddr) error
	FreeCIDR(cidr string) error
	SetStore(store Store)
	CIDR(*net.IPAddr) string
}

type GenericAlloc struct {
	store  Store
	prefix *net.IPNet
}

func NewGenericAlloc(store Store, prefix string) (Allocator, error) {
	_, pfx, err := net.ParseCIDR(prefix)
	if err != nil {
		return nil, err
	}
	return &GenericAlloc{
		store:  store,
		prefix: pfx,
	}, nil
}

func NewMemAlloc(prefix string) (Allocator, error) {
	return NewGenericAlloc(NewMemStore(), prefix)
}

func (a *GenericAlloc) Alloc() (ip *net.IPAddr, err error) {
	a.store.Lock()
	defer a.store.Unlock()
	c := ipaddr.NewCursor([]ipaddr.Prefix{*ipaddr.NewPrefix(a.prefix)})
	pos := c.Next()
	last := c.Last()
	for ; pos != nil && pos != last; pos = c.Next() {
		n := &net.IPAddr{IP: pos.IP}
		cidr := a.CIDR(n)
		ok, e := a.store.Exists(cidr)
		if e != nil {
			return nil, e
		}
		if !ok {
			a.store.Add(cidr)
			ip = n
			return
		}
	}
	err = ErrPrefixFull
	return
}

func (a *GenericAlloc) AllocCIDR() (cidr string, err error) {
	ip, err := a.Alloc()
	if err != nil {
		return
	}
	return a.CIDR(ip), nil
}

func (a *GenericAlloc) Free(ip *net.IPAddr) (err error) {
	a.store.Lock()
	defer a.store.Unlock()
	cidr := a.CIDR(ip)
	a.store.Delete(cidr)
	return
}

func (a *GenericAlloc) FreeCIDR(cidr string) (err error) {
	ip, _, err := net.ParseCIDR(cidr)
	if err != nil {
		return
	}
	return a.Free(&net.IPAddr{IP: ip})
}

func (a *GenericAlloc) CIDR(ip *net.IPAddr) string {
	return CIDR(a.prefix, ip)
}

func (a *GenericAlloc) SetStore(store Store) {
	a.store = store
}

func CIDR(prefix *net.IPNet, ip *net.IPAddr) string {
	n := new(net.IPNet)
	*n = *prefix
	n.IP = ip.IP
	return n.String()
}