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
	pscNATNets []*net.IPNet

	opt *Options
}

type Options struct {
	// If true, allocate more pscNAT CIDRs
	AllocateMoreNATs bool
}

type Option interface {
	Apply(opt *Options)
}

type allocateMoreNATs struct{}

func (allocateMoreNATs) Apply(opt *Options) {
	opt.AllocateMoreNATs = true
}

func AllocateMoreNATs() Option {
	return allocateMoreNATs{}
}

func NewGKECIDRDispatcher(cidr string, opts ...Option) (*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())
	}

	o := Options{}
	for _, opt := range opts {
		opt.Apply(&o)
	}

	return &GKECIDRDispatcher{
		prefix: prefixNet,
		opt:    &o,
	}, 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.pscNATNets = []*net.IPNet{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
}

// CutCIDR will cut a CIDR from a large CIDR
// - ipNet is the large CIDR
// - size is the size of the CIDR we want to cut
// - result is an array, the first one is the CIDR we cutted, others are left CIDRs
// e.g.
// - 10.0.0.0/16, 17
//   - 10.0.0.0/17
//   - 10.0.128.0/17
//
// - 10.0.0.0/16, 18
//   - 10.0.0.0/18
//   - 10.0.64.0/18
//   - 10.0.128.0/17
//
// - 10.0.0.0/16, 19
//   - 10.0.0.0/19
//   - 10.0.32.0/19
//   - 10.0.64.0/18
//   - 10.0.128.0/17
func CutCIDR(ipNet *net.IPNet, size int) ([]*net.IPNet, error) {
	// get start ip and mask size
	ip := ipNet.IP.To4()
	maskSize, _ := ipNet.Mask.Size()
	if maskSize > size {
		return nil, fmt.Errorf("cannot cut a CIDR larger than current one, current: %v, cut: %v", maskSize, size)
	}

	var result []*net.IPNet

	// append cutted CIDR
	result = append(result, &net.IPNet{
		IP:   ip,
		Mask: net.CIDRMask(size, 32),
	})

	for nSize := size; nSize > maskSize; nSize-- {
		result = append(result, &net.IPNet{
			IP:   nextCIDR(ip, nSize),
			Mask: net.CIDRMask(nSize, 32),
		})
	}

	return result, nil
}

// nextCIDR returns the second CIDR with newMaskSize
// It will set the last bit to 1
// e.g.
// - 10.0.0.0, 17 returns 10.0.128.0
// - 10.0.0.0, 18 returns 10.0.64.0
// - 10.0.0.0, 19 returns 10.0.32.0
// NOTE: it only worked for ipv4, and no any check for ip and mask size
func nextCIDR(ip net.IP, newMaskSize int) net.IP {
	n := len(ip)
	next := make(net.IP, n)
	copy(next, ip)
	index, bit := newMaskSize/8, newMaskSize%8
	if bit == 0 {
		index--
		bit = 8
	}
	next[index] = next[index] | (1 << (8 - bit))

	return next
}

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 {
	if c.opt.AllocateMoreNATs {
		return c.allocate()
	}

	return c.legacyAllocate()
}

// This allocate alg is deprecated, it will be removed later
// See https://pingcap.feishu.cn/wiki/HSF4woLAViypk3kCOFfcaOZOnue?from=from_copylink for more detail
func (c *GKECIDRDispatcher) legacyAllocate() 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.pscNATNets = []*net.IPNet{pscNATNet}
	return nil
}

// See https://pingcap.feishu.cn/wiki/HSF4woLAViypk3kCOFfcaOZOnue?from=from_copylink
func (c *GKECIDRDispatcher) allocate() error {
	// Assume VPC size is s
	size, _ := c.prefix.Mask.Size()
	// s -> s+1, s+1
	nets, err := CutCIDR(c.prefix, size+1)
	if err != nil {
		return err
	}

	// the size of podNet is s+1
	c.podNet = nets[0]

	// s+1 -> s+3, s+3, s+2
	nets02, err := CutCIDR(nets[1], size+3)
	if err != nil {
		return err
	}

	// the size of nodeNet is s+3
	c.nodeNet = nets02[0]
	// the size of serviceNet is s+3
	c.serviceNet = nets02[1]

	// s+2 -> s+4, s+4, s+3
	nets03, err := CutCIDR(nets02[2], size+4)
	if err != nil {
		return err
	}

	// s+4 -> 25, ..., s+5
	nets04, err := CutCIDR(nets03[0], 25)
	if err != nil {
		return err
	}
	c.pscNATNets = nets04[1:]

	// master mask is /28
	// 25 -> 28, 28, 27, 26
	nets05, err := CutCIDR(nets04[0], 28)
	if err != nil {
		return err
	}
	c.masterNet = nets05[0]

	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) PSCNATIPv4CIDRBlocks() []string {
	cidrs := make([]string, 0, len(c.pscNATNets))
	for _, net := range c.pscNATNets {
		cidrs = append(cidrs, net.String())
	}
	return cidrs
}

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
}