package util

import (
	"encoding/binary"
	"fmt"
	"net"
	"net/netip"
	"sort"
	"strings"

	"github.com/pingcap/errors"
)

const (
	aliyunPrivateLinkEndpointSubnetPrefix = 26
	aliyunOptimizedLBZoneCount            = 2
	aliyunOptimizedMgmtSplitBits          = 2
)

func GetAwsSubnetCIDRs(cidr string, zones int) (workers []string, publicLBs []string, internalLBs []string, err error) {
	prefix, err := netip.ParsePrefix(cidr)
	if err != nil {
		return nil, nil, nil, errors.AddStack(err)
	}
	zoneNames := make([]string, 0, zones)
	for i := 0; i < zones; i++ {
		zoneNames = append(zoneNames, fmt.Sprintf("%09d", i))
	}
	layout, err := BuildAWSStandardSubnetLayout(prefix.Bits(), zoneNames)
	if err != nil {
		return nil, nil, nil, err
	}
	workers, err = applySubnetShapes(prefix, layout.NodeSubnets)
	if err != nil {
		return nil, nil, nil, err
	}
	publicLBs, err = applySubnetShapes(prefix, layout.PublicLBSubnets)
	if err != nil {
		return nil, nil, nil, err
	}
	internalLBs, err = applySubnetShapes(prefix, layout.PrivateLBSubnets)
	if err != nil {
		return nil, nil, nil, err
	}
	return workers, publicLBs, internalLBs, nil
}

func GetAliyunSubnetCIDRs(cidr string, zones int) (workers []string, publicLBs []string, internalLBs []string, err error) {
	return GetAwsSubnetCIDRs(cidr, zones)
}

// GetAliyunSubnetCIDRsOptimized returns subnet CIDRs for ACK optimized node-zone mode.
//
// Layout:
//   - first 1/2 of VPC for node subnets, split by nodeZones
//   - third 1/4 for public/private LB subnets in exactly two AZs
//   - fourth 1/4 reserved for endpoint-only PrivateLink subnets
func GetAliyunSubnetCIDRsOptimized(
	cidr string,
	nodeZones []string,
	zones []string,
) (workers map[string]string, publicLBs map[string]string, internalLBs map[string]string, err error) {
	normalizedNodeZones, lbZones, err := aliyunOptimizedLBZones(nodeZones, zones)
	if err != nil {
		return nil, nil, nil, err
	}

	nodeHalf, err := getSubCIDR(cidr, 1, 0)
	if err != nil {
		return nil, nil, nil, err
	}
	mgmtHalf, err := getSubCIDR(cidr, 1, 1)
	if err != nil {
		return nil, nil, nil, err
	}
	mgmtQuarter, err := getSubCIDR(mgmtHalf, 1, 0)
	if err != nil {
		return nil, nil, nil, err
	}

	workers = make(map[string]string, len(normalizedNodeZones))
	nodeSplitBits := minSubBitsForSlots(len(normalizedNodeZones))
	for i, zone := range normalizedNodeZones {
		nodeCIDR, nodeErr := getSubCIDR(nodeHalf, nodeSplitBits, i)
		if nodeErr != nil {
			return nil, nil, nil, nodeErr
		}
		workers[zone] = nodeCIDR
	}

	publicLBs = make(map[string]string, len(lbZones))
	internalLBs = make(map[string]string, len(lbZones))
	idx := 0
	for _, zone := range lbZones {
		publicCIDR, publicErr := getSubCIDR(mgmtQuarter, aliyunOptimizedMgmtSplitBits, idx)
		if publicErr != nil {
			return nil, nil, nil, publicErr
		}
		idx++
		privateCIDR, privateErr := getSubCIDR(mgmtQuarter, aliyunOptimizedMgmtSplitBits, idx)
		if privateErr != nil {
			return nil, nil, nil, privateErr
		}
		idx++

		publicLBs[zone] = publicCIDR
		internalLBs[zone] = privateCIDR
	}

	return workers, publicLBs, internalLBs, nil
}

// GetAwsSubnetCIDRsOptimized returns subnet CIDRs for 1AZ/2AZ workload mode where node zones
// are a strict subset of all zones.
//
// Layout:
//   - first 1/2 of VPC for node subnets
//   - third 1/4 for management subnets (cp-only/public/private)
//   - fourth 1/4 reserved (unused)
func GetAwsSubnetCIDRsOptimized(
	cidr string,
	nodeZones []string,
	zones []string,
) (workers map[string]string, publicLBs map[string]string, internalLBs map[string]string, controlPlanes map[string]string, err error) {
	prefix, err := netip.ParsePrefix(cidr)
	if err != nil {
		return nil, nil, nil, nil, err
	}
	layout, err := BuildAWSOptimizedSubnetLayout(prefix.Bits(), nodeZones, zones)
	if err != nil {
		return nil, nil, nil, nil, err
	}
	workers, err = applySubnetShapeMap(prefix, layout.NodeSubnets)
	if err != nil {
		return nil, nil, nil, nil, err
	}
	publicLBs, err = applySubnetShapeMap(prefix, layout.PublicLBSubnets)
	if err != nil {
		return nil, nil, nil, nil, err
	}
	internalLBs, err = applySubnetShapeMap(prefix, layout.PrivateLBSubnets)
	if err != nil {
		return nil, nil, nil, nil, err
	}
	controlPlanes, err = applySubnetShapeMap(prefix, layout.ControlPlanes)
	if err != nil {
		return nil, nil, nil, nil, err
	}
	return workers, publicLBs, internalLBs, controlPlanes, nil
}

// GetAwsPrivateLinkEndpointSubnetCIDRs returns fixed /23 subnet CIDRs from the
// optimized layout's reserved quarter, keyed by logical AZ name. The /23 prefix
// preserves the existing EKS endpoint subnet layout.
func GetAwsPrivateLinkEndpointSubnetCIDRs(cidr string, endpointZones []string) (map[string]string, error) {
	prefix, err := netip.ParsePrefix(cidr)
	if err != nil {
		return nil, errors.AddStack(err)
	}
	layout, err := BuildAWSEndpointSubnetLayout(prefix.Bits(), endpointZones)
	if err != nil {
		return nil, err
	}
	return applySubnetShapeMap(prefix, layout.EndpointSubnets)
}

// GetAliyunPrivateLinkEndpointSubnetCIDRs returns fixed /26 vSwitch CIDRs from
// the optimized layout's reserved quarter, keyed by logical AZ name. ACK uses a
// denser prefix because Alicloud regions can expose more PrivateLink zones.
func GetAliyunPrivateLinkEndpointSubnetCIDRs(cidr string, endpointZones []string) (map[string]string, error) {
	return getPrivateLinkEndpointSubnetCIDRs(cidr, endpointZones, aliyunPrivateLinkEndpointSubnetPrefix)
}

func getPrivateLinkEndpointSubnetCIDRs(
	cidr string,
	endpointZones []string,
	subnetPrefix int,
) (map[string]string, error) {
	zones, err := normalizePrivateLinkEndpointZones(endpointZones)
	if err != nil {
		return nil, err
	}

	mgmtHalf, err := getSubCIDR(cidr, 1, 1)
	if err != nil {
		return nil, err
	}
	reservedQuarter, err := getSubCIDR(mgmtHalf, 1, 1)
	if err != nil {
		return nil, err
	}

	_, reservedNet, err := net.ParseCIDR(reservedQuarter)
	if err != nil {
		return nil, errors.AddStack(err)
	}
	ones, _ := reservedNet.Mask.Size()
	if ones > subnetPrefix {
		return nil, fmt.Errorf("reserved private link CIDR %s is too small for /%d endpoint subnets", reservedQuarter, subnetPrefix)
	}

	subBits := subnetPrefix - ones
	capacity := 1 << subBits
	if len(zones) > capacity {
		return nil, fmt.Errorf("reserved private link CIDR %s can fit %d /%d endpoint subnets, but %d zones are configured",
			reservedQuarter, capacity, subnetPrefix, len(zones))
	}

	endpoints := make(map[string]string, len(zones))
	for i, zone := range zones {
		endpointCIDR, endpointErr := getSubCIDR(reservedQuarter, subBits, i)
		if endpointErr != nil {
			return nil, endpointErr
		}
		endpoints[zone] = endpointCIDR
	}
	return endpoints, nil
}

func ShouldUsePrivateLinkEndpointSubnets(endpointZones []string, nodeZones []string) bool {
	return len(endpointZones) > 0 && len(nodeZones) == 1
}

// minSubBitsForSlots returns minimal split bits n so that 2^n >= slots.
func minSubBitsForSlots(slots int) int {
	if slots <= 1 {
		return 0
	}
	bits := 0
	capacity := 1
	for capacity < slots {
		bits++
		capacity <<= 1
	}
	return bits
}

func aliyunOptimizedLBZones(nodeZones []string, zones []string) ([]string, []string, error) {
	normalizedNodeZones, nonNodeZones, err := normalizeOptimizedZones(nodeZones, zones)
	if err != nil {
		return nil, nil, err
	}

	lbZones := make([]string, 0, aliyunOptimizedLBZoneCount)
	lbZones = append(lbZones, normalizedNodeZones...)
	for _, zone := range nonNodeZones {
		if len(lbZones) >= aliyunOptimizedLBZoneCount {
			break
		}
		lbZones = append(lbZones, zone)
	}

	return normalizedNodeZones, lbZones, nil
}

// normalizeOptimizedZones validates optimized mode constraints and returns sorted zones:
// - nodes: workload zones in stable order
// - cpOnlyZones: zones that only host control-plane related subnets
func normalizeOptimizedZones(nodeZones []string, zones []string) ([]string, []string, error) {
	if len(zones) == 0 {
		return nil, nil, fmt.Errorf("zones is empty")
	}
	if len(nodeZones) == 0 {
		return nil, nil, fmt.Errorf("node zones is empty")
	}
	if len(nodeZones) >= len(zones) {
		return nil, nil, fmt.Errorf("node zones must be a strict subset of zones")
	}

	allZones := append([]string(nil), zones...)
	sort.Strings(allZones)

	nodes := append([]string(nil), nodeZones...)
	sort.Strings(nodes)

	seenZone := make(map[string]struct{}, len(allZones))
	for _, zone := range allZones {
		if _, ok := seenZone[zone]; ok {
			return nil, nil, fmt.Errorf("zones contains duplicate zone %q", zone)
		}
		seenZone[zone] = struct{}{}
	}

	nodeSet := make(map[string]struct{}, len(nodes))
	for _, zone := range nodes {
		if _, ok := seenZone[zone]; !ok {
			return nil, nil, fmt.Errorf("node zones contains zone %q not in zones", zone)
		}
		if _, ok := nodeSet[zone]; ok {
			return nil, nil, fmt.Errorf("node zones contains duplicate zone %q", zone)
		}
		nodeSet[zone] = struct{}{}
	}

	cpOnlyZones := make([]string, 0, len(allZones)-len(nodes))
	for _, zone := range allZones {
		if _, ok := nodeSet[zone]; !ok {
			cpOnlyZones = append(cpOnlyZones, zone)
		}
	}

	if len(nodes) > 2 {
		return nil, nil, fmt.Errorf("optimized mode only supports 1 or 2 node zones")
	}

	return nodes, cpOnlyZones, nil
}

func normalizePrivateLinkEndpointZones(endpointZones []string) ([]string, error) {
	if len(endpointZones) == 0 {
		return nil, fmt.Errorf("private link endpoint zones are empty")
	}

	zones := append([]string(nil), endpointZones...)
	sort.Strings(zones)

	seen := make(map[string]struct{}, len(zones))
	for _, zone := range zones {
		trimmedZone := strings.TrimSpace(zone)
		if trimmedZone == "" {
			return nil, fmt.Errorf("private link endpoint zones contains empty zone")
		}
		if zone != trimmedZone {
			return nil, fmt.Errorf("private link endpoint zones contains untrimmed zone %q", zone)
		}
		if _, ok := seen[trimmedZone]; ok {
			return nil, fmt.Errorf("private link endpoint zones contains duplicate zone %q", trimmedZone)
		}
		seen[trimmedZone] = struct{}{}
	}
	return zones, nil
}

// getSubCIDR splits a CIDR block by subBits and returns the child at index.
// Example: getSubCIDR("10.0.0.0/16", 2, 1) => "10.0.64.0/18".
func getSubCIDR(cidr string, subBits int, index int) (string, error) {
	_, ipNet, err := net.ParseCIDR(cidr)
	if err != nil {
		return "", errors.AddStack(err)
	}

	ipInt := binary.BigEndian.Uint32(ipNet.IP)

	ones, _ := ipNet.Mask.Size()
	if ones+subBits > 32 {
		return "", fmt.Errorf("host bits is %d, but want to split by %d bits", 32-ones, subBits)
	}

	index <<= 32 - ones - subBits

	data := make([]byte, 4)
	binary.BigEndian.PutUint32(data, ipInt|uint32(index))

	sub := fmt.Sprintf("%s/%d", net.IP(data).String(), ones+subBits)
	return sub, nil
}

func applySubnetShapes(prefix netip.Prefix, shapes []AWSSubnetShape) ([]string, error) {
	cidrs := make([]string, 0, len(shapes))
	for _, shape := range shapes {
		subnet, err := shape.Apply(prefix)
		if err != nil {
			return nil, err
		}
		cidrs = append(cidrs, subnet.String())
	}
	return cidrs, nil
}

func applySubnetShapeMap(prefix netip.Prefix, shapes []AWSSubnetShape) (map[string]string, error) {
	cidrs := make(map[string]string, len(shapes))
	for _, shape := range shapes {
		subnet, err := shape.Apply(prefix)
		if err != nil {
			return nil, err
		}
		cidrs[shape.Zone] = subnet.String()
	}
	return cidrs, nil
}