package util

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

const awsEndpointSubnetPrefix = 23

type AWSSubnetShape struct {
	Zone         string
	PrefixLength int
	Index        uint64
}

type AWSSubnetLayout struct {
	NodeSubnets      []AWSSubnetShape
	PublicLBSubnets  []AWSSubnetShape
	PrivateLBSubnets []AWSSubnetShape
	ControlPlanes    []AWSSubnetShape
	EndpointSubnets  []AWSSubnetShape
}

func (s AWSSubnetShape) Apply(prefix netip.Prefix) (netip.Prefix, error) {
	if !prefix.Addr().Is4() {
		return netip.Prefix{}, fmt.Errorf("only IPv4 prefixes are supported")
	}
	root := prefix.Masked()
	rootBits := root.Bits()
	if s.PrefixLength < rootBits || s.PrefixLength > 32 {
		return netip.Prefix{}, fmt.Errorf("subnet prefix /%d is outside root prefix /%d", s.PrefixLength, rootBits)
	}
	childCount := uint64(1) << uint(s.PrefixLength-rootBits)
	if s.Index >= childCount {
		return netip.Prefix{}, fmt.Errorf("subnet index %d exceeds /%d capacity %d", s.Index, s.PrefixLength, childCount)
	}

	addrBytes := root.Addr().As4()
	base := uint64(binary.BigEndian.Uint32(addrBytes[:]))
	childSize := uint64(1) << uint(32-s.PrefixLength)
	value := base + s.Index*childSize
	if value > uint64(^uint32(0)) {
		return netip.Prefix{}, fmt.Errorf("subnet address overflows IPv4")
	}
	out := [4]byte{}
	binary.BigEndian.PutUint32(out[:], uint32(value))
	return netip.PrefixFrom(netip.AddrFrom4(out), s.PrefixLength), nil
}

func BuildAWSStandardSubnetLayout(prefixLength int, zones []string) (AWSSubnetLayout, error) {
	normalized, err := normalizeAWSZones(zones)
	if err != nil {
		return AWSSubnetLayout{}, err
	}
	if len(normalized) > 4 {
		return AWSSubnetLayout{}, fmt.Errorf("standard layout supports at most four zones")
	}
	layout := AWSSubnetLayout{
		NodeSubnets:      make([]AWSSubnetShape, 0, len(normalized)),
		PublicLBSubnets:  make([]AWSSubnetShape, 0, len(normalized)),
		PrivateLBSubnets: make([]AWSSubnetShape, 0, len(normalized)),
	}
	for i, zone := range normalized {
		layout.NodeSubnets = append(layout.NodeSubnets, AWSSubnetShape{
			Zone: zone, PrefixLength: prefixLength + 3, Index: uint64(i * 2),
		})
		layout.PublicLBSubnets = append(layout.PublicLBSubnets, AWSSubnetShape{
			Zone: zone, PrefixLength: prefixLength + 4, Index: uint64(i*4 + 2),
		})
		layout.PrivateLBSubnets = append(layout.PrivateLBSubnets, AWSSubnetShape{
			Zone: zone, PrefixLength: prefixLength + 4, Index: uint64(i*4 + 3),
		})
	}
	return layout, nil
}

func BuildAWSOptimizedSubnetLayout(prefixLength int, nodeZones, zones []string) (AWSSubnetLayout, error) {
	nodes, cpOnly, err := normalizeAWSOptimizedZones(nodeZones, zones)
	if err != nil {
		return AWSSubnetLayout{}, err
	}
	layout := AWSSubnetLayout{
		NodeSubnets:      make([]AWSSubnetShape, 0, len(nodes)),
		PublicLBSubnets:  make([]AWSSubnetShape, 0, len(nodes)),
		PrivateLBSubnets: make([]AWSSubnetShape, 0, len(nodes)),
		ControlPlanes:    make([]AWSSubnetShape, 0, len(cpOnly)),
	}
	nodeSplitBits := minAWSSubnetSplitBits(len(nodes))
	for i, zone := range nodes {
		layout.NodeSubnets = append(layout.NodeSubnets, AWSSubnetShape{
			Zone: zone, PrefixLength: prefixLength + 1 + nodeSplitBits, Index: uint64(i),
		})
	}

	requiredManagementSlots := len(cpOnly) + 2*len(nodes)
	managementSplitBits := minAWSSubnetSplitBits(requiredManagementSlots)
	managementPrefix := prefixLength + 2 + managementSplitBits
	managementBaseIndex := uint64(2 << managementSplitBits)
	index := 0
	for _, zone := range cpOnly {
		layout.ControlPlanes = append(layout.ControlPlanes, AWSSubnetShape{
			Zone: zone, PrefixLength: managementPrefix, Index: managementBaseIndex + uint64(index),
		})
		index++
	}
	for _, zone := range nodes {
		layout.PublicLBSubnets = append(layout.PublicLBSubnets, AWSSubnetShape{
			Zone: zone, PrefixLength: managementPrefix, Index: managementBaseIndex + uint64(index),
		})
		index++
		layout.PrivateLBSubnets = append(layout.PrivateLBSubnets, AWSSubnetShape{
			Zone: zone, PrefixLength: managementPrefix, Index: managementBaseIndex + uint64(index),
		})
		index++
	}
	return layout, nil
}

func BuildAWSEndpointSubnetLayout(prefixLength int, endpointZones []string) (AWSSubnetLayout, error) {
	zones, err := normalizeAWSZones(endpointZones)
	if err != nil {
		return AWSSubnetLayout{}, err
	}
	reservedPrefix := prefixLength + 2
	if reservedPrefix > awsEndpointSubnetPrefix {
		return AWSSubnetLayout{}, fmt.Errorf("reserved endpoint quarter /%d is too small for /%d endpoint subnets", reservedPrefix, awsEndpointSubnetPrefix)
	}
	splitBits := awsEndpointSubnetPrefix - reservedPrefix
	capacity := 1 << uint(splitBits)
	if len(zones) > capacity {
		return AWSSubnetLayout{}, fmt.Errorf("reserved endpoint quarter can fit %d /%d endpoint subnets, but %d zones are configured", capacity, awsEndpointSubnetPrefix, len(zones))
	}
	baseIndex := uint64(3 << splitBits)
	layout := AWSSubnetLayout{EndpointSubnets: make([]AWSSubnetShape, 0, len(zones))}
	for i, zone := range zones {
		layout.EndpointSubnets = append(layout.EndpointSubnets, AWSSubnetShape{
			Zone: zone, PrefixLength: awsEndpointSubnetPrefix, Index: baseIndex + uint64(i),
		})
	}
	return layout, nil
}

func minAWSSubnetSplitBits(slots int) int {
	if slots <= 1 {
		return 0
	}
	bits := 0
	capacity := 1
	for capacity < slots {
		bits++
		capacity <<= 1
	}
	return bits
}

func normalizeAWSZones(zones []string) ([]string, error) {
	if len(zones) == 0 {
		return nil, fmt.Errorf("zones are empty")
	}
	normalized := append([]string(nil), zones...)
	sort.Strings(normalized)
	seen := make(map[string]struct{}, len(normalized))
	for _, zone := range normalized {
		trimmed := strings.TrimSpace(zone)
		if trimmed == "" {
			return nil, fmt.Errorf("zones contains empty zone")
		}
		if zone != trimmed {
			return nil, fmt.Errorf("zones contains untrimmed zone %q", zone)
		}
		if _, ok := seen[zone]; ok {
			return nil, fmt.Errorf("zones contain duplicate value %q", zone)
		}
		seen[zone] = struct{}{}
	}
	return normalized, nil
}

func normalizeAWSOptimizedZones(nodeZones, zones []string) ([]string, []string, error) {
	allZones, err := normalizeAWSZones(zones)
	if err != nil {
		return nil, nil, err
	}
	nodes, err := normalizeAWSZones(nodeZones)
	if err != nil {
		return nil, nil, fmt.Errorf("node %w", err)
	}
	if len(nodes) >= len(allZones) {
		return nil, nil, fmt.Errorf("node zones must be a strict subset of zones")
	}
	if len(nodes) > 2 {
		return nil, nil, fmt.Errorf("optimized layout supports at most two node zones")
	}
	allSet := make(map[string]struct{}, len(allZones))
	for _, zone := range allZones {
		allSet[zone] = struct{}{}
	}
	nodeSet := make(map[string]struct{}, len(nodes))
	for _, zone := range nodes {
		if _, ok := allSet[zone]; !ok {
			return nil, nil, fmt.Errorf("node zone %q is not in region zones", zone)
		}
		nodeSet[zone] = struct{}{}
	}
	cpOnly := make([]string, 0, len(allZones)-len(nodes))
	for _, zone := range allZones {
		if _, ok := nodeSet[zone]; !ok {
			cpOnly = append(cpOnly, zone)
		}
	}
	return nodes, cpOnly, nil
}