package privatelinkgroup

import (
	"context"
	"testing"

	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/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
)

func TestEnsureVPCIPv6RetriesCIDRConflictAfterConcurrentAllocation(t *testing.T) {
	cli := &fakeVPCClient{
		vpc: ec2types.Vpc{
			VpcId: awssdk.String("vpc-1"),
			Ipv6CidrBlockAssociationSet: []ec2types.VpcIpv6CidrBlockAssociation{
				vpcAssociation("2600:1::/56", ec2types.VpcCidrBlockStateCodeAssociated),
			},
		},
		subnets: []ec2types.Subnet{{
			SubnetId: awssdk.String("subnet-target"), VpcId: awssdk.String("vpc-1"),
		}},
	}
	cli.associateSubnet = func(input *ec2.AssociateSubnetCidrBlockInput) error {
		if len(cli.associatedSubnetCIDRs) != 1 {
			return nil
		}
		cli.subnets = append(cli.subnets, ec2types.Subnet{
			SubnetId: awssdk.String("subnet-concurrent"), VpcId: awssdk.String("vpc-1"),
			Ipv6CidrBlockAssociationSet: []ec2types.SubnetIpv6CidrBlockAssociation{
				subnetAssociation(awssdk.ToString(input.Ipv6CidrBlock), ec2types.SubnetCidrBlockStateCodeAssociated),
			},
		})
		return &smithy.GenericAPIError{Code: invalidSubnetConflict}
	}

	ready, err := EnsureVPCIPv6(context.Background(), cli, "vpc-1", []string{"subnet-target"})

	require.NoError(t, err)
	assert.False(t, ready)
	assert.Equal(t, []string{"2600:1::/64", "2600:1:0:1::/64"}, cli.associatedSubnetCIDRs)
	assert.Equal(t, 2, cli.describeSubnetCalls)
}

func TestEnsureVPCIPv6AcceptsExistingSubnetCIDRFromAnotherActiveVPCPrefix(t *testing.T) {
	cli := &fakeVPCClient{
		vpc: ec2types.Vpc{
			VpcId: awssdk.String("vpc-1"),
			Ipv6CidrBlockAssociationSet: []ec2types.VpcIpv6CidrBlockAssociation{
				vpcAssociation("2600:1::/56", ec2types.VpcCidrBlockStateCodeAssociated),
				vpcAssociation("2600:2::/56", ec2types.VpcCidrBlockStateCodeAssociated),
			},
		},
		subnets: []ec2types.Subnet{{
			SubnetId: awssdk.String("subnet-1"), VpcId: awssdk.String("vpc-1"),
			Ipv6CidrBlockAssociationSet: []ec2types.SubnetIpv6CidrBlockAssociation{
				subnetAssociation("2600:2::/64", ec2types.SubnetCidrBlockStateCodeAssociated),
			},
		}},
	}

	ready, err := EnsureVPCIPv6(context.Background(), cli, "vpc-1", []string{"subnet-1"})

	require.NoError(t, err)
	assert.True(t, ready)
	assert.Empty(t, cli.associatedSubnetCIDRs)
}

func TestVPCIPv6ReadyPrefersAssociatedPrefixOverEarlierAssociatingPrefix(t *testing.T) {
	cli := &fakeVPCClient{
		vpc: ec2types.Vpc{
			VpcId: awssdk.String("vpc-1"),
			Ipv6CidrBlockAssociationSet: []ec2types.VpcIpv6CidrBlockAssociation{
				vpcAssociation("2600:1::/56", ec2types.VpcCidrBlockStateCodeAssociating),
				vpcAssociation("2600:2::/56", ec2types.VpcCidrBlockStateCodeAssociated),
			},
		},
		subnets: []ec2types.Subnet{{
			SubnetId: awssdk.String("subnet-1"), VpcId: awssdk.String("vpc-1"),
			Ipv6CidrBlockAssociationSet: []ec2types.SubnetIpv6CidrBlockAssociation{
				subnetAssociation("2600:2::/64", ec2types.SubnetCidrBlockStateCodeAssociated),
			},
		}},
	}

	ready, err := VPCIPv6Ready(context.Background(), cli, "vpc-1", []string{"subnet-1"})

	require.NoError(t, err)
	assert.True(t, ready)
}

func TestEnsureVPCIPv6PrefersAssociatedSubnetCIDROverEarlierAssociatingCIDR(t *testing.T) {
	cli := &fakeVPCClient{
		vpc: ec2types.Vpc{
			VpcId: awssdk.String("vpc-1"),
			Ipv6CidrBlockAssociationSet: []ec2types.VpcIpv6CidrBlockAssociation{
				vpcAssociation("2600:1::/56", ec2types.VpcCidrBlockStateCodeAssociated),
			},
		},
		subnets: []ec2types.Subnet{{
			SubnetId: awssdk.String("subnet-1"), VpcId: awssdk.String("vpc-1"),
			Ipv6CidrBlockAssociationSet: []ec2types.SubnetIpv6CidrBlockAssociation{
				subnetAssociation("2600:0::/64", ec2types.SubnetCidrBlockStateCodeAssociating),
				subnetAssociation("2600:1::/64", ec2types.SubnetCidrBlockStateCodeAssociated),
			},
		}},
	}

	ready, err := EnsureVPCIPv6(context.Background(), cli, "vpc-1", []string{"subnet-1"})

	require.NoError(t, err)
	assert.True(t, ready)
	assert.Empty(t, cli.associatedSubnetCIDRs)
}

func vpcAssociation(cidr string, state ec2types.VpcCidrBlockStateCode) ec2types.VpcIpv6CidrBlockAssociation {
	return ec2types.VpcIpv6CidrBlockAssociation{
		Ipv6CidrBlock:      awssdk.String(cidr),
		Ipv6CidrBlockState: &ec2types.VpcCidrBlockState{State: state},
	}
}

func subnetAssociation(cidr string, state ec2types.SubnetCidrBlockStateCode) ec2types.SubnetIpv6CidrBlockAssociation {
	return ec2types.SubnetIpv6CidrBlockAssociation{
		Ipv6CidrBlock:      awssdk.String(cidr),
		Ipv6CidrBlockState: &ec2types.SubnetCidrBlockState{State: state},
	}
}

type fakeVPCClient struct {
	vpc                   ec2types.Vpc
	subnets               []ec2types.Subnet
	associatedSubnetCIDRs []string
	associateSubnet       func(*ec2.AssociateSubnetCidrBlockInput) error
	describeVPCCalls      int
	describeSubnetCalls   int
}

func (f *fakeVPCClient) DescribeVpcs(
	context.Context, *ec2.DescribeVpcsInput, ...func(*ec2.Options),
) (*ec2.DescribeVpcsOutput, error) {
	f.describeVPCCalls++
	return &ec2.DescribeVpcsOutput{Vpcs: []ec2types.Vpc{f.vpc}}, nil
}

func (*fakeVPCClient) AssociateVpcCidrBlock(
	context.Context, *ec2.AssociateVpcCidrBlockInput, ...func(*ec2.Options),
) (*ec2.AssociateVpcCidrBlockOutput, error) {
	return &ec2.AssociateVpcCidrBlockOutput{}, nil
}

func (f *fakeVPCClient) DescribeSubnets(
	context.Context, *ec2.DescribeSubnetsInput, ...func(*ec2.Options),
) (*ec2.DescribeSubnetsOutput, error) {
	f.describeSubnetCalls++
	return &ec2.DescribeSubnetsOutput{Subnets: f.subnets}, nil
}

func (f *fakeVPCClient) AssociateSubnetCidrBlock(
	_ context.Context, input *ec2.AssociateSubnetCidrBlockInput, _ ...func(*ec2.Options),
) (*ec2.AssociateSubnetCidrBlockOutput, error) {
	f.associatedSubnetCIDRs = append(f.associatedSubnetCIDRs, awssdk.ToString(input.Ipv6CidrBlock))
	if f.associateSubnet != nil {
		if err := f.associateSubnet(input); err != nil {
			return nil, err
		}
	}
	return &ec2.AssociateSubnetCidrBlockOutput{}, nil
}