package privatelinkgroup

import (
	"context"
	stderrors "errors"
	"net/netip"
	"sort"

	awssdk "github.com/aws/aws-sdk-go-v2/aws"
	"github.com/aws/aws-sdk-go-v2/service/ec2"
	ec2types "github.com/aws/aws-sdk-go-v2/service/ec2/types"
	"github.com/aws/smithy-go"
	"github.com/pingcap/errors"
)

const invalidSubnetConflict = "InvalidSubnet.Conflict"

type VPCClient interface {
	DescribeVpcs(context.Context, *ec2.DescribeVpcsInput, ...func(*ec2.Options)) (*ec2.DescribeVpcsOutput, error)
	AssociateVpcCidrBlock(context.Context, *ec2.AssociateVpcCidrBlockInput, ...func(*ec2.Options)) (*ec2.AssociateVpcCidrBlockOutput, error)
	DescribeSubnets(context.Context, *ec2.DescribeSubnetsInput, ...func(*ec2.Options)) (*ec2.DescribeSubnetsOutput, error)
	AssociateSubnetCidrBlock(context.Context, *ec2.AssociateSubnetCidrBlockInput, ...func(*ec2.Options)) (*ec2.AssociateSubnetCidrBlockOutput, error)
}

// EnsureVPCIPv6 monotonically adds one VPC /56 and one /64 only to the
// subnets actually attached to the group NLB. It never removes IPv6 CIDRs.
func EnsureVPCIPv6(ctx context.Context, cli VPCClient, vpcID string, subnetIDs []string) (bool, error) {
	vpc, err := describeVPC(ctx, cli, vpcID)
	if err != nil {
		return false, err
	}
	vpcCIDR, ready, err := ensureVPCIPv6CIDR(ctx, cli, vpc)
	if err != nil || !ready {
		return ready, err
	}
	return ensureSubnetIPv6CIDRs(ctx, cli, vpcID, vpcCIDR, subnetIDs)
}

// VPCIPv6Ready is the read-only counterpart used by API evaluation.
func VPCIPv6Ready(ctx context.Context, cli VPCClient, vpcID string, subnetIDs []string) (bool, error) {
	if len(subnetIDs) == 0 {
		return false, errors.New("PrivateLink NLB attached subnet set is empty")
	}
	vpc, err := describeVPC(ctx, cli, vpcID)
	if err != nil {
		return false, err
	}
	_, ready := observedVPCIPv6CIDR(vpc)
	if !ready {
		return false, nil
	}
	subnets, err := describeSubnets(ctx, cli, vpcID)
	if err != nil {
		return false, err
	}
	byID := make(map[string]ec2types.Subnet, len(subnets))
	for _, subnet := range subnets {
		byID[awssdk.ToString(subnet.SubnetId)] = subnet
	}
	for _, subnetID := range subnetIDs {
		subnet, ok := byID[subnetID]
		if !ok || awssdk.ToString(subnet.VpcId) != vpcID {
			return false, errors.Errorf("PrivateLink NLB subnet %s is not in VPC %s", subnetID, vpcID)
		}
		ready := false
		for _, association := range subnet.Ipv6CidrBlockAssociationSet {
			if association.Ipv6CidrBlockState != nil &&
				association.Ipv6CidrBlockState.State == ec2types.SubnetCidrBlockStateCodeAssociated &&
				validSubnetIPv6CIDR(awssdk.ToString(association.Ipv6CidrBlock)) {
				ready = true
				break
			}
		}
		if !ready {
			return false, nil
		}
	}
	return true, nil
}

func describeVPC(ctx context.Context, cli VPCClient, vpcID string) (ec2types.Vpc, error) {
	if cli == nil {
		return ec2types.Vpc{}, errors.New("AWS VPC client is not initialized")
	}
	out, err := cli.DescribeVpcs(ctx, &ec2.DescribeVpcsInput{VpcIds: []string{vpcID}})
	if err != nil {
		return ec2types.Vpc{}, errors.Wrapf(err, "failed to describe VPC %s", vpcID)
	}
	if out == nil || len(out.Vpcs) == 0 {
		return ec2types.Vpc{}, errors.Errorf("AWS did not return VPC %s", vpcID)
	}
	return out.Vpcs[0], nil
}

func observedVPCIPv6CIDR(vpc ec2types.Vpc) (string, bool) {
	active := activeVPCIPv6Associations(vpc.Ipv6CidrBlockAssociationSet)
	sort.Slice(active, func(i, j int) bool {
		return awssdk.ToString(active[i].Ipv6CidrBlock) < awssdk.ToString(active[j].Ipv6CidrBlock)
	})
	var associating string
	for i := range active {
		cidr := awssdk.ToString(active[i].Ipv6CidrBlock)
		if _, err := vpcIPv6Prefix(cidr); err != nil {
			continue
		}
		if active[i].Ipv6CidrBlockState.State == ec2types.VpcCidrBlockStateCodeAssociated {
			return cidr, true
		}
		if associating == "" {
			associating = cidr
		}
	}
	if associating != "" {
		return associating, false
	}
	return "", false
}

func ensureVPCIPv6CIDR(ctx context.Context, cli VPCClient, vpc ec2types.Vpc) (string, bool, error) {
	cidr, ready := observedVPCIPv6CIDR(vpc)
	if cidr != "" {
		return cidr, ready, nil
	}
	_, err := cli.AssociateVpcCidrBlock(ctx, &ec2.AssociateVpcCidrBlockInput{
		VpcId: vpc.VpcId, AmazonProvidedIpv6CidrBlock: awssdk.Bool(true),
	})
	return "", false, err
}

func ensureSubnetIPv6CIDRs(ctx context.Context, cli VPCClient, vpcID, vpcCIDR string, subnetIDs []string) (bool, error) {
	if len(subnetIDs) == 0 {
		return false, errors.New("PrivateLink NLB attached subnet set is empty")
	}
	prefix, err := vpcIPv6Prefix(vpcCIDR)
	if err != nil {
		return false, err
	}
	subnets, err := describeSubnets(ctx, cli, vpcID)
	if err != nil {
		return false, err
	}
	byID := make(map[string]ec2types.Subnet, len(subnets))
	used := make(map[string]struct{})
	for _, subnet := range subnets {
		byID[awssdk.ToString(subnet.SubnetId)] = subnet
		for _, association := range subnet.Ipv6CidrBlockAssociationSet {
			cidr := awssdk.ToString(association.Ipv6CidrBlock)
			if subnetIPv6CIDRInPrefix(cidr, prefix) {
				used[cidr] = struct{}{}
			}
		}
	}
	targets := append([]string(nil), subnetIDs...)
	sort.Strings(targets)
	ready := true
	for _, subnetID := range targets {
		subnet, ok := byID[subnetID]
		if !ok || awssdk.ToString(subnet.VpcId) != vpcID {
			return false, errors.Errorf("PrivateLink NLB subnet %s is not in VPC %s", subnetID, vpcID)
		}
		hasAssociated, hasAssociating := observedSubnetIPv6State(subnet)
		if hasAssociated {
			continue
		}
		if hasAssociating {
			ready = false
			continue
		}
		for {
			cidr, err := firstUnusedSubnetIPv6CIDR(prefix, used)
			if err != nil {
				return false, err
			}
			if _, err := cli.AssociateSubnetCidrBlock(ctx, &ec2.AssociateSubnetCidrBlockInput{
				SubnetId: awssdk.String(subnetID), Ipv6CidrBlock: awssdk.String(cidr),
			}); err == nil {
				used[cidr] = struct{}{}
				ready = false
				break
			} else if !isInvalidSubnetConflict(err) {
				return false, errors.Wrapf(err, "failed to associate %s with subnet %s", cidr, subnetID)
			}

			// Another Group transition in the same VPC can select the same free
			// /64 from an identical DescribeSubnets snapshot. Refresh AWS state
			// after the conflict: if this subnet was updated concurrently, wait
			// for that association; otherwise exclude the contested /64 and try
			// the next one. This keeps allocation safe across processes without
			// relying on an in-memory lock.
			used[cidr] = struct{}{}
			refreshed, refreshErr := describeSubnets(ctx, cli, vpcID)
			if refreshErr != nil {
				return false, refreshErr
			}
			for _, current := range refreshed {
				byID[awssdk.ToString(current.SubnetId)] = current
				for _, association := range current.Ipv6CidrBlockAssociationSet {
					currentCIDR := awssdk.ToString(association.Ipv6CidrBlock)
					if subnetIPv6CIDRInPrefix(currentCIDR, prefix) {
						used[currentCIDR] = struct{}{}
					}
				}
			}
			subnet, ok = byID[subnetID]
			if !ok || awssdk.ToString(subnet.VpcId) != vpcID {
				return false, errors.Errorf("PrivateLink NLB subnet %s is not in VPC %s", subnetID, vpcID)
			}
			hasAssociated, hasAssociating = observedSubnetIPv6State(subnet)
			if hasAssociated {
				break
			}
			if hasAssociating {
				ready = false
				break
			}
		}
	}
	return ready, nil
}

func observedSubnetIPv6State(subnet ec2types.Subnet) (hasAssociated, hasAssociating bool) {
	for _, association := range activeSubnetIPv6Associations(subnet.Ipv6CidrBlockAssociationSet) {
		if !validSubnetIPv6CIDR(awssdk.ToString(association.Ipv6CidrBlock)) {
			continue
		}
		if association.Ipv6CidrBlockState.State == ec2types.SubnetCidrBlockStateCodeAssociated {
			hasAssociated = true
			continue
		}
		hasAssociating = true
	}
	return hasAssociated, hasAssociating
}

func isInvalidSubnetConflict(err error) bool {
	var apiErr smithy.APIError
	return stderrors.As(err, &apiErr) && apiErr.ErrorCode() == invalidSubnetConflict
}

func describeSubnets(ctx context.Context, cli VPCClient, vpcID string) ([]ec2types.Subnet, error) {
	var result []ec2types.Subnet
	var token *string
	for {
		out, err := cli.DescribeSubnets(ctx, &ec2.DescribeSubnetsInput{
			Filters:    []ec2types.Filter{{Name: awssdk.String("vpc-id"), Values: []string{vpcID}}},
			MaxResults: awssdk.Int32(1000), NextToken: token,
		})
		if err != nil {
			return nil, errors.Wrapf(err, "failed to describe VPC %s subnets", vpcID)
		}
		if out == nil {
			return nil, errors.Errorf("AWS returned an empty subnet page for VPC %s", vpcID)
		}
		result = append(result, out.Subnets...)
		token = out.NextToken
		if awssdk.ToString(token) == "" {
			return result, nil
		}
	}
}

func vpcIPv6Prefix(cidr string) (netip.Prefix, error) {
	prefix, err := netip.ParsePrefix(cidr)
	if err != nil {
		return netip.Prefix{}, errors.Wrapf(err, "invalid VPC IPv6 CIDR %q", cidr)
	}
	prefix = prefix.Masked()
	if !prefix.Addr().Is6() || prefix.Bits() != 56 {
		return netip.Prefix{}, errors.Errorf("expected an IPv6 /56, got %s", prefix)
	}
	return prefix, nil
}

func subnetIPv6CIDRInPrefix(cidr string, vpcPrefix netip.Prefix) bool {
	prefix, err := netip.ParsePrefix(cidr)
	if err != nil {
		return false
	}
	prefix = prefix.Masked()
	return prefix.Addr().Is6() && prefix.Bits() == 64 && vpcPrefix.Contains(prefix.Addr())
}

func validSubnetIPv6CIDR(cidr string) bool {
	prefix, err := netip.ParsePrefix(cidr)
	if err != nil {
		return false
	}
	prefix = prefix.Masked()
	return prefix.Addr().Is6() && prefix.Bits() == 64
}

func firstUnusedSubnetIPv6CIDR(vpcPrefix netip.Prefix, used map[string]struct{}) (string, error) {
	base := vpcPrefix.Addr().As16()
	for slot := 0; slot < 256; slot++ {
		candidate := base
		// netip.Addr.As16 returns a fixed [16]byte value, so index 7 is always valid.
		//nolint:gosec // G602 is a false positive for this fixed-size array.
		candidate[7] = byte(slot)
		cidr := netip.PrefixFrom(netip.AddrFrom16(candidate), 64).String()
		if _, exists := used[cidr]; !exists {
			return cidr, nil
		}
	}
	return "", errors.Errorf("no unused IPv6 /64 remains in %s", vpcPrefix)
}

func activeVPCIPv6Associations(values []ec2types.VpcIpv6CidrBlockAssociation) []ec2types.VpcIpv6CidrBlockAssociation {
	result := make([]ec2types.VpcIpv6CidrBlockAssociation, 0, len(values))
	for _, value := range values {
		if value.Ipv6CidrBlockState != nil && (value.Ipv6CidrBlockState.State == ec2types.VpcCidrBlockStateCodeAssociating || value.Ipv6CidrBlockState.State == ec2types.VpcCidrBlockStateCodeAssociated) {
			result = append(result, value)
		}
	}
	return result
}

func activeSubnetIPv6Associations(values []ec2types.SubnetIpv6CidrBlockAssociation) []ec2types.SubnetIpv6CidrBlockAssociation {
	result := make([]ec2types.SubnetIpv6CidrBlockAssociation, 0, len(values))
	for _, value := range values {
		if value.Ipv6CidrBlockState != nil && (value.Ipv6CidrBlockState.State == ec2types.SubnetCidrBlockStateCodeAssociating || value.Ipv6CidrBlockState.State == ec2types.SubnetCidrBlockStateCodeAssociated) {
			result = append(result, value)
		}
	}
	return result
}