package cidr

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

// GKECIDRDispatcher dispatcher CIDR with the last 12 bits(prefix /20), including
// - Pod IPv4 CIDR block(0, /21)
// - Node IPv4 CIDR block(100, /23)
// - Service IPv4 CIDR block(110, /23)
// - Master IPv4 CIDR block(11100000, /28)
type GKECIDRDispatcher struct {
	prefix *net.IPNet

	podNet, nodeNet, serviceNet, masterNet *net.IPNet

	// CIDR for PSC NAT
	pscNATNet *net.IPNet
}

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

	if size, _ := prefixNet.Mask.Size(); size > 20 {
		return nil, fmt.Errorf("too large CIDR prefix, size: %d", size)
	}
	if len(prefixNet.IP.To4()) != 4 {
		return nil, fmt.Errorf("not an IPv4 address, %s", prefixNet.IP.To4().String())
	}

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

// Default will only used for migration cluster.
func (c *GKECIDRDispatcher) Default() {
	// Pod
	podNet := c.prefixDeepCopy()
	podNet.Mask = net.CIDRMask(21, 32)
	c.podNet = podNet

	// Node
	nodeNet := c.prefixDeepCopy()
	nodeNet.IP.To4()[2] |= 0b00001000
	nodeNet.Mask = net.CIDRMask(23, 32)
	c.nodeNet = nodeNet

	// Service
	serviceNet := c.prefixDeepCopy()
	serviceNet.IP.To4()[2] |= 0b00001100
	serviceNet.Mask = net.CIDRMask(23, 32)
	c.serviceNet = serviceNet

	// Master
	masterNet := c.prefixDeepCopy()
	masterNet.IP.To4()[2] |= 0b00001110
	masterNet.Mask = net.CIDRMask(28, 32)
	c.masterNet = masterNet

	// PSC NAT
	pscNATNet := c.prefixDeepCopy()
	pscNATNet.IP.To4()[2] |= 0b00001110
	pscNATNet.IP.To4()[3] |= 0b10000000
	pscNATNet.Mask = net.CIDRMask(25, 32)
	c.pscNATNet = pscNATNet
}

func SplitSubnet(ipNet *net.IPNet, count int) ([]*net.IPNet, error) {
	// get start ip and mask size
	ip := ipNet.IP.To4()
	maskSize, _ := ipNet.Mask.Size()

	// caculate new mask size
	newMaskSize := maskSize + int(math.Log2(float64(count)))

	// check if new mask size is valid
	if newMaskSize > 32 {
		return nil, fmt.Errorf("can not split subnets %d", count)
	}

	// create subnets list
	subnets := make([]*net.IPNet, count)

	// caculate subnets size
	subnetSize := int(math.Pow(2, float64(32-newMaskSize)))

	// allocate subnets
	for i := 0; i < count; i++ {
		start := make(net.IP, len(ip))
		copy(start, ip)
		startInt := IPToInt(start)
		startInt += uint32(i) * uint32(subnetSize)

		subnet := &net.IPNet{
			IP:   IntToIP(startInt),
			Mask: net.CIDRMask(newMaskSize, 32),
		}

		subnets[i] = subnet
	}

	return subnets, nil
}

func IPToInt(ip net.IP) uint32 {
	return binary.BigEndian.Uint32(ip.To4())
}

func IntToIP(i uint32) net.IP {
	ip := make(net.IP, 4)
	binary.BigEndian.PutUint32(ip, i)
	return ip
}

func (c *GKECIDRDispatcher) Allocate() error {
	// split subnets to 2 parts
	nets, err := SplitSubnet(c.prefix, 2)
	if err != nil {
		return err
	}

	// part 0 for pod and part 1 for node and service
	c.podNet = nets[0]

	nets02, err := SplitSubnet(nets[1], 4)
	if err != nil {
		return err
	}

	c.nodeNet = nets02[0]
	c.serviceNet = nets02[1]

	ones, _ := nets02[2].Mask.Size()
	// master mask is /28
	nets03, err := SplitSubnet(nets02[2], int(math.Pow(2, float64(28-ones))))
	if err != nil {
		return err
	}

	c.masterNet = nets03[0]

	// PSC NAT
	pscNATNet := deepCopy(nets03[0])
	pscNATNet.IP.To4()[3] |= 0b10000000
	pscNATNet.Mask = net.CIDRMask(25, 32)
	c.pscNATNet = pscNATNet
	return nil
}

func deepCopy(old *net.IPNet) *net.IPNet {
	cidr := new(net.IPNet)

	cidr.IP = make([]byte, len(old.IP))
	copy(cidr.IP, old.IP)

	cidr.Mask = make([]byte, len(old.Mask))
	copy(cidr.Mask, old.Mask)

	return cidr
}

func (c *GKECIDRDispatcher) GetPrefixNetDeepCopy() *net.IPNet {
	return c.prefixDeepCopy()
}

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

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

func (c *GKECIDRDispatcher) ServiceIPv4CIDRBlock() string {
	return c.serviceNet.String()
}

func (c *GKECIDRDispatcher) MasterIPv4CIDRBlock() string {
	return c.masterNet.String()
}

func (c *GKECIDRDispatcher) PSCNATIPv4CIDRBlock() string {
	return c.pscNATNet.String()
}

func (c *GKECIDRDispatcher) prefixDeepCopy() *net.IPNet {
	cidr := new(net.IPNet)

	cidr.IP = make([]byte, len(c.prefix.IP))
	copy(cidr.IP, c.prefix.IP)

	cidr.Mask = make([]byte, len(c.prefix.Mask))
	copy(cidr.Mask, c.prefix.Mask)

	return cidr
}

// The 20th bit is reserved when the network cidr prefix size is 19
// Set it to 1 to avoid CIDR overlap
func (c *GKECIDRDispatcher) SetPrefixPlaceholder() error {
	if ones, _ := c.prefix.Mask.Size(); ones > 19 {
		return fmt.Errorf("prefix size is %d, bigger then 19, no placeholder left", ones)
	}
	c.prefix.IP.To4()[2] |= 0b00010000
	c.Default()
	return nil
}
