package awgserver

import (
	"fmt"
	"net/netip"
)

// IPAllocator assigns IPs from a CIDR subnet for WireGuard peers.
type IPAllocator struct {
	prefix netip.Prefix
}

// NewIPAllocator creates an allocator for the given CIDR (e.g., "10.14.0.0/24").
func NewIPAllocator(cidr string) (*IPAllocator, error) {
	prefix, err := netip.ParsePrefix(cidr)
	if err != nil {
		return nil, fmt.Errorf("parse CIDR %q: %w", cidr, err)
	}
	prefix = prefix.Masked()
	return &IPAllocator{prefix: prefix}, nil
}

// ServerIP returns the first usable IP in the subnet (network + 1), used as the server address.
func (a *IPAllocator) ServerIP() netip.Addr {
	return a.prefix.Addr().Next()
}

// Allocate finds the next free IP in the subnet, skipping the network address,
// server address (.1), broadcast, and any IPs in the used set.
func (a *IPAllocator) Allocate(used []string) (string, error) {
	usedSet := make(map[netip.Addr]struct{}, len(used))
	for _, u := range used {
		p, err := netip.ParsePrefix(u)
		if err != nil {
			continue
		}
		usedSet[p.Addr()] = struct{}{}
	}

	// Start from network + 2 (skip network addr and server addr).
	addr := a.prefix.Addr().Next().Next()
	for a.prefix.Contains(addr) {
		next := addr.Next()
		// Skip if next would be outside prefix (broadcast equivalent).
		if !a.prefix.Contains(next) {
			break
		}
		if _, taken := usedSet[addr]; !taken {
			return fmt.Sprintf("%s/32", addr.String()), nil
		}
		addr = next
	}

	return "", fmt.Errorf("no free IPs in %s", a.prefix.String())
}