package cidr

import (
	"encoding/binary"
	"fmt"
	"net"
)

// aksCIDRDispatcher is the dispatcher for AKS CIDR.
type aksCIDRDispatcher struct {
	// prefix is the CIDR prefix of the VPC.
	prefix *net.IPNet
	// podNet is the CIDR block for AKS pod.
	// By default, different nodepools in the same aks use the same pod subnet.
	podNet *net.IPNet
	// nodeNet is the CIDR block for AKS node.
	// By default, different nodepools in the same aks use the same node subnet.
	nodeNet *net.IPNet
	// serviceNet is the CIDR block for AKS service.
	serivceNet *net.IPNet
	// plsNATNet is the CIDR block for AKS PrivateLink Service NAT.
	// By default, different PrivateLink Service in the same aks use the same pls nat subnet.
	plsNATNet *net.IPNet
	// apiServerNet is the delegated subnet for AKS API server VNet integration.
	apiServerNet *net.IPNet
	// dnsIP is the DNS IP for the AKS.
	// By default, the DNS IP is the 11th IP of the service CIDR.
	dnsIP net.IP
}

func NewAKSCIDRDispatcher(cidr string) (*aksCIDRDispatcher, error) {
	_, prefixNet, err := net.ParseCIDR(cidr)
	if err != nil {
		return nil, err
	}

	if len(prefixNet.IP.To4()) != 4 {
		return nil, fmt.Errorf("not an IPv4 address, %s", prefixNet.IP.To4().String())
	}

	return &aksCIDRDispatcher{
		prefix: prefixNet,
	}, nil
}

// Allocate will allocate the CIDR blocks for AKS.
// Assume the CIDR prefix is /19, the CIDR blocks are:
// - Pod IPv4 CIDR block(00000, /20)
// - Node IPv4 CIDR block(10000, /21)
// - Service IPv4 CIDR block(11000, /22)
// - PLS NAT IPv4 CIDR block(11100, /23)
// - API server delegated subnet (/28) from the reserved block
// - Reserved for future use
func (c *aksCIDRDispatcher) Allocate() error {
	// split the subnet into 2 parts
	nets, err := SplitSubnet(c.prefix, 2)
	if err != nil {
		return err
	}
	// pard 0 for pod subnet, part 1 for node subne/pls nat subnet/service cidr
	c.podNet = nets[0]

	// split the subnet into 2 parts
	nets02, err := SplitSubnet(nets[1], 2)
	if err != nil {
		return err
	}
	// part 0 of nets02 for node subnet, part 1 for service cidr/pls nat subnet
	c.nodeNet = nets02[0]

	nets03, err := SplitSubnet(nets02[1], 2)
	if err != nil {
		return err
	}
	// part 0 of nets03 for service cidr, part 1 for pls nat subnet
	c.serivceNet = nets03[0]

	// get the dns ip from the service cidr
	// the dns ip is the 11th ip of the service cidr
	dnsIP, err := getDNSIPFromCIDR(c.serivceNet.String(), 10)
	if err != nil {
		return err
	}
	c.dnsIP = dnsIP

	nets04, err := SplitSubnet(nets03[1], 2)
	if err != nil {
		return err
	}
	// part 0 of nets04 for pls nat subnet, part 1 of nets04 reserved for future use
	c.plsNATNet = nets04[0]
	apiServerNet, err := firstSubnetWithMask(nets04[1], 28)
	if err != nil {
		return err
	}
	c.apiServerNet = apiServerNet

	return nil
}

func (c *aksCIDRDispatcher) PodIPv4CIDRBlock() string {
	return c.podNet.String()
}

func (c *aksCIDRDispatcher) NodeIPv4CIDRBlock() string {
	return c.nodeNet.String()
}

func (c *aksCIDRDispatcher) ServiceIPv4CIDRBlock() string {
	return c.serivceNet.String()
}

func (c *aksCIDRDispatcher) PLSNATIPv4CIDRBlock() string {
	return c.plsNATNet.String()
}

func (c *aksCIDRDispatcher) APIServerIPv4CIDRBlock() string {
	return c.apiServerNet.String()
}

func (c *aksCIDRDispatcher) DNSServiceIP() string {
	return c.dnsIP.String()
}

func getDNSIPFromCIDR(cidr string, idx uint32) (net.IP, error) {
	_, ipnet, err := net.ParseCIDR(cidr)
	if err != nil {
		return nil, err
	}
	network := ipnet.IP.To4()
	ipInt := binary.BigEndian.Uint32(network)
	ipInt += idx
	dnsIP := make(net.IP, 4)
	binary.BigEndian.PutUint32(dnsIP, ipInt)

	if !ipnet.Contains(dnsIP) {
		return nil, fmt.Errorf("dns ip %s is not in the cidr %s", dnsIP.String(), cidr)
	}

	return dnsIP, nil
}

func firstSubnetWithMask(parent *net.IPNet, mask int) (*net.IPNet, error) {
	if parent == nil {
		return nil, fmt.Errorf("parent subnet is nil")
	}
	parentMask, _ := parent.Mask.Size()
	if mask < parentMask || mask > 32 {
		return nil, fmt.Errorf("invalid subnet mask /%d for parent %s", mask, parent.String())
	}
	return &net.IPNet{
		IP:   append(net.IP(nil), parent.IP.To4()...),
		Mask: net.CIDRMask(mask, 32),
	}, nil
}