package dataplane

import (
	"context"
	"errors"
	"testing"

	"github.com/stretchr/testify/require"
	"google.golang.org/grpc/codes"
	"google.golang.org/grpc/status"

	dmpbv1 "github.com/tidbcloud/cluster-service-ng/api/tidbcloud/dataplane-manager/v1"
	rsconfig "github.com/tidbcloud/infra-provider/regional-server/pkg/config"
	"github.com/tidbcloud/infra-provider/regional-server/pkg/consts"
)

type fakeEndpointZoneResolver struct {
	zones      []string
	err        error
	calls      int
	roleARN    string
	externalID string
	region     string
}

func (f *fakeEndpointZoneResolver) Resolve(
	_ context.Context,
	roleARN string,
	externalID string,
	region string,
) ([]string, error) {
	f.calls++
	f.roleARN = roleARN
	f.externalID = externalID
	f.region = region
	return f.zones, f.err
}

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

	t.Run("regional AWS capacity does not discover endpoint zones", func(t *testing.T) {
		t.Parallel()
		resolver := &fakeEndpointZoneResolver{err: errors.New("must not be called")}
		service := newNetworkCapacityService(consts.ProviderTypeEKS, resolver)

		response, err := service.EvaluateDataPlaneNetworkCapacity(t.Context(), capacityRequest(23, dmpbv1.DataPlane_REGIONAL))

		require.NoError(t, err)
		require.True(t, response.GetSupported())
		require.EqualValues(t, 31, response.GetSupportedNodeCount())
		require.Zero(t, resolver.calls)
	})

	t.Run("zonal AWS capacity discovers endpoint zones", func(t *testing.T) {
		t.Parallel()
		resolver := &fakeEndpointZoneResolver{zones: []string{"us-west-2c", "us-west-2a", "us-west-2b"}}
		service := newNetworkCapacityService(consts.ProviderTypeEKS, resolver)
		request := capacityRequest(19, dmpbv1.DataPlane_ZONAL)
		request.AccountId = " 123456789012 "

		response, err := service.EvaluateDataPlaneNetworkCapacity(t.Context(), request)

		require.NoError(t, err)
		require.True(t, response.GetSupported())
		require.Positive(t, response.GetSupportedNodeCount())
		require.Equal(t, 1, resolver.calls)
		require.Equal(t, "arn:aws:iam::123456789012:role/DataPlaneRole", resolver.roleARN)
		require.Equal(t, "external-id", resolver.externalID)
		require.Equal(t, "us-west-2", resolver.region)
	})

	t.Run("explicit zonal node zone works without a configured default", func(t *testing.T) {
		t.Parallel()
		resolver := &fakeEndpointZoneResolver{zones: []string{"us-west-2a", "us-west-2b", "us-west-2c"}}
		service := newNetworkCapacityService(consts.ProviderTypeEKS, resolver)
		service.config.RegionConfig.NodeZones = nil
		request := capacityRequest(19, dmpbv1.DataPlane_ZONAL)
		request.NodeZones = []string{"us-west-2b"}
		request.AccountId = "123456789012"

		response, err := service.EvaluateDataPlaneNetworkCapacity(t.Context(), request)

		require.NoError(t, err)
		require.True(t, response.GetSupported())
		require.Positive(t, response.GetSupportedNodeCount())
		require.Equal(t, 1, resolver.calls)
	})

	t.Run("regional node zones are accepted but deployment zones determine capacity", func(t *testing.T) {
		t.Parallel()
		service := newNetworkCapacityService(consts.ProviderTypeEKS, nil)
		request := capacityRequest(23, dmpbv1.DataPlane_REGIONAL)
		request.NodeZones = []string{"ignored-a", "ignored-b"}

		response, err := service.EvaluateDataPlaneNetworkCapacity(t.Context(), request)

		require.NoError(t, err)
		require.True(t, response.GetSupported())
		require.EqualValues(t, 31, response.GetSupportedNodeCount())
	})

	t.Run("zonal BYOC capacity does not discover endpoint zones", func(t *testing.T) {
		t.Parallel()
		resolver := &fakeEndpointZoneResolver{err: errors.New("must not be called")}
		service := newNetworkCapacityService(consts.ProviderTypeEKS, resolver)
		service.config.BYOC.BYOCPrivateLink.ControlPlaneVPCID = "vpc-control-plane"

		response, err := service.EvaluateDataPlaneNetworkCapacity(
			t.Context(),
			capacityRequest(19, dmpbv1.DataPlane_ZONAL),
		)

		require.NoError(t, err)
		require.True(t, response.GetSupported())
		require.Zero(t, resolver.calls)
	})

	t.Run("physical endpoint capacity failure returns an OK response", func(t *testing.T) {
		t.Parallel()
		resolver := &fakeEndpointZoneResolver{zones: []string{"us-west-2a", "us-west-2b", "us-west-2c"}}
		service := newNetworkCapacityService(consts.ProviderTypeEKS, resolver)
		request := capacityRequest(20, dmpbv1.DataPlane_ZONAL)
		request.AccountId = "123456789012"

		response, err := service.EvaluateDataPlaneNetworkCapacity(t.Context(), request)

		require.NoError(t, err)
		require.False(t, response.GetSupported())
		require.Zero(t, response.GetSupportedNodeCount())
	})
}

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

	tests := []struct {
		name    string
		service *DataPlaneService
		request *dmpbv1.EvaluateDataPlaneNetworkCapacityRequest
		code    codes.Code
	}{
		{
			name:    "missing prefix",
			service: newNetworkCapacityService(consts.ProviderTypeEKS, nil),
			request: &dmpbv1.EvaluateDataPlaneNetworkCapacityRequest{HighAvailabilityType: dmpbv1.DataPlane_REGIONAL},
			code:    codes.InvalidArgument,
		},
		{
			name:    "prefix below IPv4 bound",
			service: newNetworkCapacityService(consts.ProviderTypeEKS, nil),
			request: capacityRequest(-1, dmpbv1.DataPlane_REGIONAL),
			code:    codes.InvalidArgument,
		},
		{
			name:    "prefix above IPv4 bound",
			service: newNetworkCapacityService(consts.ProviderTypeEKS, nil),
			request: capacityRequest(33, dmpbv1.DataPlane_REGIONAL),
			code:    codes.InvalidArgument,
		},
		{
			name:    "unspecified availability",
			service: newNetworkCapacityService(consts.ProviderTypeEKS, nil),
			request: capacityRequest(20, dmpbv1.DataPlane_HIGH_AVAILABILITY_TYPE_UNSPECIFIED),
			code:    codes.InvalidArgument,
		},
		{
			name:    "unknown availability",
			service: newNetworkCapacityService(consts.ProviderTypeEKS, nil),
			request: capacityRequest(20, dmpbv1.DataPlane_HighAvailabilityType(99)),
			code:    codes.InvalidArgument,
		},
		{
			name:    "unsupported provider",
			service: newNetworkCapacityService(consts.ProviderTypeACK, nil),
			request: capacityRequest(20, dmpbv1.DataPlane_REGIONAL),
			code:    codes.FailedPrecondition,
		},
		{
			name:    "node zone on regional evaluation",
			service: newNetworkCapacityService(consts.ProviderTypeEKS, nil),
			request: func() *dmpbv1.EvaluateDataPlaneNetworkCapacityRequest {
				req := capacityRequest(20, dmpbv1.DataPlane_REGIONAL)
				req.NodeZones = []string{"us-west-2a"}
				return req
			}(),
			code: codes.InvalidArgument,
		},
		{
			name:    "empty node zone",
			service: newNetworkCapacityService(consts.ProviderTypeEKS, nil),
			request: func() *dmpbv1.EvaluateDataPlaneNetworkCapacityRequest {
				req := capacityRequest(20, dmpbv1.DataPlane_ZONAL)
				req.NodeZones = []string{""}
				return req
			}(),
			code: codes.InvalidArgument,
		},
		{
			name:    "node zone outside configured region zones",
			service: newNetworkCapacityService(consts.ProviderTypeEKS, nil),
			request: func() *dmpbv1.EvaluateDataPlaneNetworkCapacityRequest {
				req := capacityRequest(20, dmpbv1.DataPlane_ZONAL)
				req.NodeZones = []string{"us-east-1a"}
				return req
			}(),
			code: codes.InvalidArgument,
		},
		{
			name:    "multiple node zones on zonal evaluation",
			service: newNetworkCapacityService(consts.ProviderTypeEKS, nil),
			request: func() *dmpbv1.EvaluateDataPlaneNetworkCapacityRequest {
				req := capacityRequest(20, dmpbv1.DataPlane_ZONAL)
				req.NodeZones = []string{"us-west-2a", "us-west-2b"}
				return req
			}(),
			code: codes.InvalidArgument,
		},
		{
			name:    "zonal endpoint discovery needs an account",
			service: newNetworkCapacityService(consts.ProviderTypeEKS, &fakeEndpointZoneResolver{}),
			request: capacityRequest(19, dmpbv1.DataPlane_ZONAL),
			code:    codes.FailedPrecondition,
		},
		{
			name:    "zonal endpoint discovery needs a resolver",
			service: newNetworkCapacityService(consts.ProviderTypeEKS, nil),
			request: func() *dmpbv1.EvaluateDataPlaneNetworkCapacityRequest {
				req := capacityRequest(19, dmpbv1.DataPlane_ZONAL)
				req.AccountId = "123456789012"
				return req
			}(),
			code: codes.Internal,
		},
		{
			name: "endpoint discovery failure",
			service: newNetworkCapacityService(
				consts.ProviderTypeEKS,
				&fakeEndpointZoneResolver{err: errors.New("AWS unavailable")},
			),
			request: func() *dmpbv1.EvaluateDataPlaneNetworkCapacityRequest {
				req := capacityRequest(19, dmpbv1.DataPlane_ZONAL)
				req.AccountId = "123456789012"
				return req
			}(),
			code: codes.Unavailable,
		},
		{
			name: "endpoint discovery deadline",
			service: newNetworkCapacityService(
				consts.ProviderTypeEKS,
				&fakeEndpointZoneResolver{err: context.DeadlineExceeded},
			),
			request: func() *dmpbv1.EvaluateDataPlaneNetworkCapacityRequest {
				req := capacityRequest(19, dmpbv1.DataPlane_ZONAL)
				req.AccountId = "123456789012"
				return req
			}(),
			code: codes.DeadlineExceeded,
		},
		{
			name:    "empty endpoint discovery result",
			service: newNetworkCapacityService(consts.ProviderTypeEKS, &fakeEndpointZoneResolver{}),
			request: func() *dmpbv1.EvaluateDataPlaneNetworkCapacityRequest {
				req := capacityRequest(19, dmpbv1.DataPlane_ZONAL)
				req.AccountId = "123456789012"
				return req
			}(),
			code: codes.Unavailable,
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			t.Parallel()
			response, err := tt.service.EvaluateDataPlaneNetworkCapacity(t.Context(), tt.request)
			require.Nil(t, response)
			require.Equal(t, tt.code, status.Code(err))
		})
	}

	t.Run("unsupported provider is rejected before endpoint configuration", func(t *testing.T) {
		t.Parallel()
		service := newNetworkCapacityService(consts.ProviderTypeACK, nil)
		_, err := service.EvaluateDataPlaneNetworkCapacity(t.Context(), capacityRequest(20, dmpbv1.DataPlane_REGIONAL))
		require.Equal(t, codes.FailedPrecondition, status.Code(err))
		require.ErrorContains(t, err, `unsupported for provider "ack"`)
	})

	t.Run("endpoint discovery needs a configured role", func(t *testing.T) {
		t.Parallel()
		service := newNetworkCapacityService(consts.ProviderTypeEKS, &fakeEndpointZoneResolver{})
		service.config.DataPlaneRoleName = ""
		request := capacityRequest(19, dmpbv1.DataPlane_ZONAL)
		request.AccountId = "123456789012"
		_, err := service.EvaluateDataPlaneNetworkCapacity(t.Context(), request)
		require.Equal(t, codes.FailedPrecondition, status.Code(err))
	})

	t.Run("zonal evaluation needs one configured node zone", func(t *testing.T) {
		t.Parallel()
		service := newNetworkCapacityService(consts.ProviderTypeEKS, nil)
		service.config.RegionConfig.NodeZones = nil
		_, err := service.EvaluateDataPlaneNetworkCapacity(
			t.Context(),
			capacityRequest(20, dmpbv1.DataPlane_ZONAL),
		)
		require.Equal(t, codes.FailedPrecondition, status.Code(err))
	})
}

func newNetworkCapacityService(
	provider string,
	resolver *fakeEndpointZoneResolver,
) *DataPlaneService {
	reservedAdminPodIPs := int64(50)
	service := &DataPlaneService{
		config: rsconfig.RootConfig{
			ControllerConfig: rsconfig.ControllerConfig{
				Provider: provider,
				RegionConfig: rsconfig.RegionConfig{
					Name:      "us-west-2",
					Zones:     []string{"us-west-2a", "us-west-2b", "us-west-2c"},
					NodeZones: []string{"us-west-2a"},
				},
				DataPlaneRoleName:      "DataPlaneRole",
				AWSDataPlaneExternalID: "external-id",
			},
			NetworkCapacity: rsconfig.NetworkCapacityConfig{ReservedAdminPodIPs: &reservedAdminPodIPs},
		},
	}
	if resolver != nil {
		service.endpointZoneResolver = resolver
	}
	return service
}

func capacityRequest(
	prefix int32,
	highAvailability dmpbv1.DataPlane_HighAvailabilityType,
) *dmpbv1.EvaluateDataPlaneNetworkCapacityRequest {
	return &dmpbv1.EvaluateDataPlaneNetworkCapacityRequest{
		CidrPrefixLength:     &prefix,
		HighAvailabilityType: highAvailability,
	}
}