package alloc

import (
	"context"
	"fmt"
	"net"

	"git.giftfish.de/ston1th/haproxy-lb/pkg/cluster"
	"github.com/mikioh/ipaddr"
)

type Alloc struct {
	kv      cluster.KV
	pool    []*net.IPNet
	gateway *net.IPNet
}

func NewAlloc(kv cluster.KV, cidrs []string, gateway string) (*Alloc, error) {
	pool, err := parseRange(cidrs)
	if err != nil {
		return nil, err
	}
	_, gw, err := net.ParseCIDR(gateway)
	if err != nil {
		return nil, err
	}
	return &Alloc{kv, pool, gw}, nil
}

func (a *Alloc) AllocIP(ctx context.Context, name string) (addr string, err error) {
	for _, cidr := range a.pool {
		c := ipaddr.NewCursor([]ipaddr.Prefix{*ipaddr.NewPrefix(a.gateway)})
		for pos := c.First(); pos != nil; pos = c.Next() {
			ip := pos.IP.String()
			v, _ := a.kv.Get(ip)
			vs := string(v)
			if v == "" {
				err = a.kv.Set(ip, []byte(name))
				if err != nil {
					return
				}
				err = a.kv.Set(name, []byte(ip))
				if err != nil {
					a.kv.Delete(ip)
					return
				}
				var gw *net.IPNet
				*gw = *a.gateway
				gw.IP = pos.IP
				addr = gw.String()
				return
			}
		}
	}
	err = fmt.Errorf("no available IPs in pool")
	return
}

func (a *Alloc) FreeIP(ctx context.Context, name string) (err error) {
	if name == "" {
		return
	}
	v, err := a.kv.Get(name)
	if err != nil {
		return
	}
	vs := string(v)
	err = a.kv.Delete(vs)
	if err != nil {
		return
	}
	err = a.kv.Delete(name)
	if err != nil {
		a.kv.Set(vs, []byte(name))
	}
	return
}

func parseRange(cidrs []string) (nets []*net.IPNet, err error) {
	for _, cidr := range cidrs {
		if !strings.Contains(cidr, "-") {
			_, n, err := net.ParseCIDR(cidr)
			if err != nil {
				return nil, err
			}
			nets = append(nets, n)
			continue
		}

		fs := strings.SplitN(cidr, "-", 2)
		if len(fs) != 2 {
			return nil, fmt.Errorf("invalid IP range %q", cidr)
		}
		start := net.ParseIP(strings.TrimSpace(fs[0]))
		if start == nil {
			return nil, fmt.Errorf("invalid IP range %q: invalid start IP %q", cidr, fs[0])
		}
		end := net.ParseIP(strings.TrimSpace(fs[1]))
		if end == nil {
			return nil, fmt.Errorf("invalid IP range %q: invalid end IP %q", cidr, fs[1])
		}

		for _, pfx := range ipaddr.Summarize(start, end) {
			n := &net.IPNet{
				IP:   pfx.IP,
				Mask: pfx.Mask,
			}
			nets = append(nets, n)
		}
	}
	return
}