package util

import (
	"net/netip"
	"testing"

	"github.com/stretchr/testify/require"
)

func TestBuildAWSStandardSubnetLayout(t *testing.T) {
	t.Parallel()

	layout, err := BuildAWSStandardSubnetLayout(20, []string{"us-west-2c", "us-west-2a", "us-west-2b"})
	require.NoError(t, err)
	require.Equal(t, []AWSSubnetShape{
		{Zone: "us-west-2a", PrefixLength: 23, Index: 0},
		{Zone: "us-west-2b", PrefixLength: 23, Index: 2},
		{Zone: "us-west-2c", PrefixLength: 23, Index: 4},
	}, layout.NodeSubnets)
	require.Equal(t, []AWSSubnetShape{
		{Zone: "us-west-2a", PrefixLength: 24, Index: 2},
		{Zone: "us-west-2b", PrefixLength: 24, Index: 6},
		{Zone: "us-west-2c", PrefixLength: 24, Index: 10},
	}, layout.PublicLBSubnets)
	require.Equal(t, []AWSSubnetShape{
		{Zone: "us-west-2a", PrefixLength: 24, Index: 3},
		{Zone: "us-west-2b", PrefixLength: 24, Index: 7},
		{Zone: "us-west-2c", PrefixLength: 24, Index: 11},
	}, layout.PrivateLBSubnets)

	_, err = BuildAWSStandardSubnetLayout(20, []string{"a", "b", "c", "d", "e"})
	require.ErrorContains(t, err, "at most four zones")
}

func TestBuildAWSEndpointSubnetLayout(t *testing.T) {
	t.Parallel()

	layout, err := BuildAWSEndpointSubnetLayout(19, []string{"us-west-2c", "us-west-2a", "us-west-2b"})
	require.NoError(t, err)
	require.Equal(t, []AWSSubnetShape{
		{Zone: "us-west-2a", PrefixLength: 23, Index: 12},
		{Zone: "us-west-2b", PrefixLength: 23, Index: 13},
		{Zone: "us-west-2c", PrefixLength: 23, Index: 14},
	}, layout.EndpointSubnets)

	_, err = BuildAWSEndpointSubnetLayout(20, []string{"a", "b", "c"})
	require.ErrorContains(t, err, "can fit 2")
	_, err = BuildAWSEndpointSubnetLayout(19, []string{"a", "a"})
	require.ErrorContains(t, err, "duplicate")
}

func TestAWSSubnetLayoutShapeAppliesDeterministically(t *testing.T) {
	t.Parallel()

	layout, err := BuildAWSOptimizedSubnetLayout(
		18,
		[]string{"us-west-2b", "us-west-2a"},
		[]string{"us-west-2c", "us-west-2a", "us-west-2b"},
	)
	require.NoError(t, err)
	require.Equal(t, []AWSSubnetShape{
		{Zone: "us-west-2a", PrefixLength: 20, Index: 0},
		{Zone: "us-west-2b", PrefixLength: 20, Index: 1},
	}, layout.NodeSubnets)

	first, err := layout.NodeSubnets[1].Apply(netip.MustParsePrefix("10.0.0.0/18"))
	require.NoError(t, err)
	second, err := layout.NodeSubnets[1].Apply(netip.MustParsePrefix("10.1.0.0/18"))
	require.NoError(t, err)
	require.Equal(t, "10.0.16.0/20", first.String())
	require.Equal(t, "10.1.16.0/20", second.String())
	require.Equal(t, first.Bits(), second.Bits())
}

func TestAWSSubnetShapeApplyRejectsInvalidInput(t *testing.T) {
	t.Parallel()

	tests := []struct {
		name    string
		prefix  netip.Prefix
		shape   AWSSubnetShape
		message string
	}{
		{
			name:    "IPv6 root",
			prefix:  netip.MustParsePrefix("2001:db8::/64"),
			shape:   AWSSubnetShape{PrefixLength: 64},
			message: "only IPv4 prefixes are supported",
		},
		{
			name:    "subnet prefix shorter than root",
			prefix:  netip.MustParsePrefix("10.0.0.0/20"),
			shape:   AWSSubnetShape{PrefixLength: 19},
			message: "outside root prefix",
		},
		{
			name:    "subnet prefix exceeds IPv4",
			prefix:  netip.MustParsePrefix("10.0.0.0/20"),
			shape:   AWSSubnetShape{PrefixLength: 33},
			message: "outside root prefix",
		},
		{
			name:    "subnet index exceeds root",
			prefix:  netip.MustParsePrefix("10.0.0.0/20"),
			shape:   AWSSubnetShape{PrefixLength: 22, Index: 4},
			message: "exceeds /22 capacity 4",
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			t.Parallel()

			_, err := tt.shape.Apply(tt.prefix)
			require.ErrorContains(t, err, tt.message)
		})
	}
}