package util

import (
	"testing"

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

func TestGetSubCIDR(t *testing.T) {
	testcases := []struct {
		cidr           string
		subBits        int
		index          int
		expectedResult string
		err            bool
	}{
		{
			cidr: "abcd",
			err:  true,
		},
		{
			cidr: "192.168.0.0",
			err:  true,
		},
		{
			cidr:    "192.168.0.0/24",
			subBits: 10,
			err:     true,
		},
		{
			cidr:           "192.168.0.0/24",
			subBits:        0,
			index:          0,
			expectedResult: "192.168.0.0/24",
		},
		{
			cidr:           "192.168.0.0/24",
			subBits:        4,
			index:          0,
			expectedResult: "192.168.0.0/28",
		},
		{
			cidr:           "192.168.0.0/24",
			subBits:        4,
			index:          1,
			expectedResult: "192.168.0.16/28",
		},
		{
			cidr:           "192.168.0.0/24",
			subBits:        4,
			index:          2,
			expectedResult: "192.168.0.32/28",
		},
		{
			cidr:           "192.168.0.0/24",
			subBits:        4,
			index:          3,
			expectedResult: "192.168.0.48/28",
		},
	}
	for _, tc := range testcases {
		subCidr, err := getSubCIDR(tc.cidr, tc.subBits, tc.index)
		if tc.err {
			assert.Error(t, err)
		} else {
			assert.NoError(t, err)
			assert.Equal(t, tc.expectedResult, subCidr)
		}
	}
}

func TestGetAwsSubnetCIDRs(t *testing.T) {
	t.Run("invalid CIDR", func(t *testing.T) {
		_, _, _, err := GetAwsSubnetCIDRs("192.168.0.0", 3)
		assert.Error(t, err)
	})
	t.Run("valid CIDR", func(t *testing.T) {
		workers, publics, privates, err := GetAwsSubnetCIDRs("192.168.0.0/16", 3)
		assert.NoError(t, err)

		expectedWorkers := []string{"192.168.0.0/19", "192.168.64.0/19", "192.168.128.0/19"}
		expectedPublics := []string{"192.168.32.0/20", "192.168.96.0/20", "192.168.160.0/20"}
		expectedPrivates := []string{"192.168.48.0/20", "192.168.112.0/20", "192.168.176.0/20"}
		assert.Equal(t, expectedWorkers, workers)
		assert.Equal(t, expectedPublics, publics)
		assert.Equal(t, expectedPrivates, privates)
	})
	t.Run("rejects unsupported layout inputs", func(t *testing.T) {
		_, _, _, err := GetAwsSubnetCIDRs("10.0.0.0/16", 0)
		assert.ErrorContains(t, err, "zones are empty")

		_, _, _, err = GetAwsSubnetCIDRs("10.0.0.0/16", 5)
		assert.ErrorContains(t, err, "at most four zones")

		_, _, _, err = GetAwsSubnetCIDRs("2001:db8::/56", 3)
		assert.ErrorContains(t, err, "only IPv4 prefixes are supported")
	})
}

func TestGetAliyunSubnetCIDRsMatchesStandardLayout(t *testing.T) {
	workers, publics, privates, err := GetAliyunSubnetCIDRs("192.168.0.0/16", 3)
	assert.NoError(t, err)
	assert.Equal(t, []string{"192.168.0.0/19", "192.168.64.0/19", "192.168.128.0/19"}, workers)
	assert.Equal(t, []string{"192.168.32.0/20", "192.168.96.0/20", "192.168.160.0/20"}, publics)
	assert.Equal(t, []string{"192.168.48.0/20", "192.168.112.0/20", "192.168.176.0/20"}, privates)
}

func TestGetAwsSubnetCIDRsOptimized(t *testing.T) {
	t.Run("1AZ workload with 2 zones", func(t *testing.T) {
		workers, publics, privates, cps, err := GetAwsSubnetCIDRsOptimized(
			"10.0.0.0/18",
			[]string{"us-west-2a"},
			[]string{"us-west-2a", "us-west-2b"},
		)
		assert.NoError(t, err)
		assert.Equal(t, map[string]string{"us-west-2a": "10.0.0.0/19"}, workers)
		assert.Equal(t, map[string]string{"us-west-2a": "10.0.36.0/22"}, publics)
		assert.Equal(t, map[string]string{"us-west-2a": "10.0.40.0/22"}, privates)
		assert.Equal(t, map[string]string{"us-west-2b": "10.0.32.0/22"}, cps)
	})

	t.Run("2AZ workload with 3 zones", func(t *testing.T) {
		workers, publics, privates, cps, err := GetAwsSubnetCIDRsOptimized(
			"10.0.0.0/18",
			[]string{"us-west-2b", "us-west-2a"},
			[]string{"us-west-2c", "us-west-2a", "us-west-2b"},
		)
		assert.NoError(t, err)
		assert.Equal(t, map[string]string{
			"us-west-2a": "10.0.0.0/20",
			"us-west-2b": "10.0.16.0/20",
		}, workers)
		assert.Equal(t, map[string]string{
			"us-west-2a": "10.0.34.0/23",
			"us-west-2b": "10.0.38.0/23",
		}, publics)
		assert.Equal(t, map[string]string{
			"us-west-2a": "10.0.36.0/23",
			"us-west-2b": "10.0.40.0/23",
		}, privates)
		assert.Equal(t, map[string]string{"us-west-2c": "10.0.32.0/23"}, cps)
	})

	t.Run("invalid node zones", func(t *testing.T) {
		_, _, _, _, err := GetAwsSubnetCIDRsOptimized("10.0.0.0/18", nil, []string{"a", "b"})
		assert.Error(t, err)

		_, _, _, _, err = GetAwsSubnetCIDRsOptimized("10.0.0.0/18", []string{"a", "a"}, []string{"a", "b"})
		assert.Error(t, err)

		_, _, _, _, err = GetAwsSubnetCIDRsOptimized("10.0.0.0/18", []string{"a", "b"}, []string{"a", "b"})
		assert.Error(t, err)

		_, _, _, _, err = GetAwsSubnetCIDRsOptimized("10.0.0.0/18", []string{"a", "b", "c"}, []string{"a", "b", "c", "d"})
		assert.Error(t, err)
	})
}

func TestGetAliyunSubnetCIDRsOptimized(t *testing.T) {
	t.Run("1AZ workload with 2AZ LB subnets", func(t *testing.T) {
		workers, publics, privates, err := GetAliyunSubnetCIDRsOptimized(
			"10.0.0.0/16",
			[]string{"ap-southeast-1a"},
			[]string{"ap-southeast-1a", "ap-southeast-1b", "ap-southeast-1c"},
		)
		assert.NoError(t, err)

		assert.Equal(t, map[string]string{
			"ap-southeast-1a": "10.0.0.0/17",
		}, workers)
		assert.Equal(t, map[string]string{
			"ap-southeast-1a": "10.0.128.0/20",
			"ap-southeast-1b": "10.0.160.0/20",
		}, publics)
		assert.Equal(t, map[string]string{
			"ap-southeast-1a": "10.0.144.0/20",
			"ap-southeast-1b": "10.0.176.0/20",
		}, privates)
	})

	t.Run("2AZ workload with 2AZ LB subnets", func(t *testing.T) {
		workers, publics, privates, err := GetAliyunSubnetCIDRsOptimized(
			"10.0.0.0/16",
			[]string{"ap-southeast-1b", "ap-southeast-1a"},
			[]string{"ap-southeast-1c", "ap-southeast-1a", "ap-southeast-1b"},
		)
		assert.NoError(t, err)

		assert.Equal(t, map[string]string{
			"ap-southeast-1a": "10.0.0.0/18",
			"ap-southeast-1b": "10.0.64.0/18",
		}, workers)
		assert.Equal(t, map[string]string{
			"ap-southeast-1a": "10.0.128.0/20",
			"ap-southeast-1b": "10.0.160.0/20",
		}, publics)
		assert.Equal(t, map[string]string{
			"ap-southeast-1a": "10.0.144.0/20",
			"ap-southeast-1b": "10.0.176.0/20",
		}, privates)
	})
}

func TestGetAwsPrivateLinkEndpointSubnetCIDRs(t *testing.T) {
	t.Run("/18 supports 6 endpoint zones", func(t *testing.T) {
		endpoints, err := GetAwsPrivateLinkEndpointSubnetCIDRs(
			"10.0.0.0/18",
			[]string{"us-east-1f", "us-east-1a", "us-east-1c", "us-east-1b", "us-east-1e", "us-east-1d"},
		)
		assert.NoError(t, err)
		assert.Equal(t, map[string]string{
			"us-east-1a": "10.0.48.0/23",
			"us-east-1b": "10.0.50.0/23",
			"us-east-1c": "10.0.52.0/23",
			"us-east-1d": "10.0.54.0/23",
			"us-east-1e": "10.0.56.0/23",
			"us-east-1f": "10.0.58.0/23",
		}, endpoints)
	})

	t.Run("/19 does not support 6 endpoint zones", func(t *testing.T) {
		_, err := GetAwsPrivateLinkEndpointSubnetCIDRs(
			"10.0.0.0/19",
			[]string{"us-east-1a", "us-east-1b", "us-east-1c", "us-east-1d", "us-east-1e", "us-east-1f"},
		)
		assert.Error(t, err)
		assert.Contains(t, err.Error(), "can fit 4 /23 endpoint subnets")
	})

	t.Run("invalid endpoint zones", func(t *testing.T) {
		_, err := GetAwsPrivateLinkEndpointSubnetCIDRs("10.0.0.0/18", nil)
		assert.Error(t, err)

		_, err = GetAwsPrivateLinkEndpointSubnetCIDRs("10.0.0.0/18", []string{"us-east-1a", "us-east-1a"})
		assert.Error(t, err)

		_, err = GetAwsPrivateLinkEndpointSubnetCIDRs("10.0.0.0/18", []string{"us-east-1a "})
		assert.Error(t, err)
		assert.Contains(t, err.Error(), "untrimmed zone")
	})
}

func TestGetAliyunPrivateLinkEndpointSubnetCIDRs(t *testing.T) {
	t.Run("allocates deterministic subnets", func(t *testing.T) {
		endpoints, err := GetAliyunPrivateLinkEndpointSubnetCIDRs(
			"10.0.0.0/18",
			[]string{
				"cn-hangzhou-k",
				"cn-hangzhou-b",
				"cn-hangzhou-j",
				"cn-hangzhou-c",
				"cn-hangzhou-i",
				"cn-hangzhou-d",
				"cn-hangzhou-h",
				"cn-hangzhou-e",
				"cn-hangzhou-g",
				"cn-hangzhou-f",
			},
		)
		require.NoError(t, err)
		assert.Equal(t, map[string]string{
			"cn-hangzhou-b": "10.0.48.0/26",
			"cn-hangzhou-c": "10.0.48.64/26",
			"cn-hangzhou-d": "10.0.48.128/26",
			"cn-hangzhou-e": "10.0.48.192/26",
			"cn-hangzhou-f": "10.0.49.0/26",
			"cn-hangzhou-g": "10.0.49.64/26",
			"cn-hangzhou-h": "10.0.49.128/26",
			"cn-hangzhou-i": "10.0.49.192/26",
			"cn-hangzhou-j": "10.0.50.0/26",
			"cn-hangzhou-k": "10.0.50.64/26",
		}, endpoints)
	})

	t.Run("rejects zones beyond reserved quarter capacity", func(t *testing.T) {
		_, err := GetAliyunPrivateLinkEndpointSubnetCIDRs(
			"10.0.0.0/21",
			[]string{"a", "b", "c", "d", "e", "f", "g", "h", "i"},
		)
		require.Error(t, err)
		assert.ErrorContains(t, err, "can fit 8 /26 endpoint subnets")
	})

	t.Run("rejects reserved quarter smaller than an endpoint subnet", func(t *testing.T) {
		_, err := GetAliyunPrivateLinkEndpointSubnetCIDRs("10.0.0.0/25", []string{"a"})
		require.Error(t, err)
		assert.ErrorContains(t, err, "is too small for /26 endpoint subnets")
	})
}