package resourcepool

import (
	"context"
	"testing"

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

	globalcommonv1 "github.com/tidbcloud/tidb-management-service/api/spec/global/common/v1"
	globalngv1 "github.com/tidbcloud/tidb-management-service/api/spec/global/nextgen/v1"
)

func TestGetDefaultDemand_Azure(t *testing.T) {
	demand := getDefaultDemand(context.Background(), globalcommonv1.CloudProvider_AZURE)
	require.NotNil(t, demand)

	assert.EqualValues(t, 20000, demand.GetInitialRcu())
	require.NotNil(t, demand.GetTikvData())
	require.NotNil(t, demand.GetTikvRaftlog())
	require.NotNil(t, demand.GetTiflashCn())
	require.NotNil(t, demand.GetTiflashWn())

	assert.Equal(t, globalngv1.ResourcePoolProfile_PremiumSSDV2, demand.GetTikvData().GetStorageType())
	assert.EqualValues(t, 1500, demand.GetTikvData().GetDiskSizeGi())
	assert.EqualValues(t, 8000, demand.GetTikvData().GetDiskIops())
	assert.EqualValues(t, 400, demand.GetTikvData().GetDiskThroughputMbps())

	assert.Equal(t, globalngv1.ResourcePoolProfile_PremiumSSDV2, demand.GetTikvRaftlog().GetStorageType())
	assert.EqualValues(t, 40, demand.GetTikvRaftlog().GetDiskSizeGi())
	assert.EqualValues(t, 3000, demand.GetTikvRaftlog().GetDiskIops())
	assert.EqualValues(t, 125, demand.GetTikvRaftlog().GetDiskThroughputMbps())

	assert.Equal(t, globalngv1.ResourcePoolProfile_PremiumSSDV2, demand.GetTiflashCn().GetStorageType())
	assert.EqualValues(t, 500, demand.GetTiflashCn().GetDiskSizeGi())
	assert.EqualValues(t, 6000, demand.GetTiflashCn().GetDiskIops())
	assert.EqualValues(t, 625, demand.GetTiflashCn().GetDiskThroughputMbps())

	assert.Equal(t, globalngv1.ResourcePoolProfile_PremiumSSDV2, demand.GetTiflashWn().GetStorageType())
	assert.EqualValues(t, 500, demand.GetTiflashWn().GetDiskSizeGi())
	assert.EqualValues(t, 6000, demand.GetTiflashWn().GetDiskIops())
	assert.EqualValues(t, 625, demand.GetTiflashWn().GetDiskThroughputMbps())
}

func TestGetDefaultDemand_UnsupportedProvider(t *testing.T) {
	demand := getDefaultDemand(context.Background(), globalcommonv1.CloudProvider_NAME_UNSPECIFIED)
	require.NotNil(t, demand)
	assert.EqualValues(t, 0, demand.GetInitialRcu())
	assert.Nil(t, demand.GetTikvData())
}

func TestFillingCreateRequestWithDefault_RcuDrivenEnvDefaults(t *testing.T) {
	tests := []struct {
		name             string
		env              string
		initialRcu       uint32
		wantInitialRcu   uint32
		wantDataSize     uint32
		wantDataIops     uint32
		wantDataTput     uint32
		wantRaftlogIops  uint32
		wantRaftlogTput  uint32
		wantTiflashIops  uint32
		wantTiflashTput  uint32
		wantStorageType  globalngv1.ResourcePoolProfile_StorageType
		wantTiflashSize  uint32
		wantRaftlogSize  uint32
		wantInstanceLmt  uint64
		wantDefaultCIDR  string
		wantDefaultSCIDR string
	}{
		{
			name:             "dev uses small disk defaults",
			env:              "dev",
			wantInitialRcu:   defaultInitialRcu,
			wantDataSize:     100,
			wantDataIops:     3000,
			wantDataTput:     125,
			wantRaftlogIops:  3000,
			wantRaftlogTput:  125,
			wantTiflashIops:  3000,
			wantTiflashTput:  125,
			wantStorageType:  globalngv1.ResourcePoolProfile_GP3,
			wantTiflashSize:  500,
			wantRaftlogSize:  40,
			wantInstanceLmt:  50,
			wantDefaultCIDR:  "10.0.0.0/16",
			wantDefaultSCIDR: "172.20.0.0/16",
		},
		{
			name:             "staging uses small disk defaults",
			env:              "staging",
			initialRcu:       60000,
			wantInitialRcu:   60000,
			wantDataSize:     100,
			wantDataIops:     3000,
			wantDataTput:     125,
			wantRaftlogIops:  3000,
			wantRaftlogTput:  125,
			wantTiflashIops:  3000,
			wantTiflashTput:  125,
			wantStorageType:  globalngv1.ResourcePoolProfile_GP3,
			wantTiflashSize:  500,
			wantRaftlogSize:  40,
			wantInstanceLmt:  50,
			wantDefaultCIDR:  "10.0.0.0/16",
			wantDefaultSCIDR: "172.20.0.0/16",
		},
		{
			name:             "prod below boundary uses first disk tier",
			env:              "prod",
			initialRcu:       49999,
			wantInitialRcu:   49999,
			wantDataSize:     500,
			wantDataIops:     3000,
			wantDataTput:     125,
			wantRaftlogIops:  3000,
			wantRaftlogTput:  125,
			wantTiflashIops:  3000,
			wantTiflashTput:  125,
			wantStorageType:  globalngv1.ResourcePoolProfile_GP3,
			wantTiflashSize:  500,
			wantRaftlogSize:  40,
			wantInstanceLmt:  50,
			wantDefaultCIDR:  "10.0.0.0/16",
			wantDefaultSCIDR: "172.20.0.0/16",
		},
		{
			name:             "prod boundary uses next disk tier",
			env:              "prod",
			initialRcu:       50000,
			wantInitialRcu:   50000,
			wantDataSize:     1000,
			wantDataIops:     4000,
			wantDataTput:     300,
			wantRaftlogIops:  4000,
			wantRaftlogTput:  125,
			wantTiflashIops:  3000,
			wantTiflashTput:  125,
			wantStorageType:  globalngv1.ResourcePoolProfile_GP3,
			wantTiflashSize:  500,
			wantRaftlogSize:  40,
			wantInstanceLmt:  50,
			wantDefaultCIDR:  "10.0.0.0/16",
			wantDefaultSCIDR: "172.20.0.0/16",
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			request := newCreateResourcePoolRequest(globalcommonv1.CloudProvider_AWS)
			request.Demand = &globalngv1.CreateResourcePoolRequest_Demand{
				InitialRcu: tt.initialRcu,
			}

			fillingCreateRequestWithDefault(context.Background(), request, tt.env, true)

			assert.Equal(t, tt.wantInitialRcu, request.GetDemand().GetInitialRcu())
			assert.Equal(t, tt.wantDefaultCIDR, request.GetResourcePool().GetCidr())
			assert.Equal(t, tt.wantDefaultSCIDR, request.GetResourcePool().GetK8SServiceCidr())
			assert.Equal(t, tt.wantInstanceLmt, request.GetResourcePool().GetInstanceLimit())
			assertStorage(t, request.GetDemand().GetTikvData(), tt.wantStorageType, tt.wantDataSize, tt.wantDataIops, tt.wantDataTput)
			assertStorage(t, request.GetDemand().GetTikvRaftlog(), tt.wantStorageType, tt.wantRaftlogSize, tt.wantRaftlogIops, tt.wantRaftlogTput)
			assertStorage(t, request.GetDemand().GetTiflashCn(), tt.wantStorageType, tt.wantTiflashSize, tt.wantTiflashIops, tt.wantTiflashTput)
			assertStorage(t, request.GetDemand().GetTiflashWn(), tt.wantStorageType, tt.wantTiflashSize, tt.wantTiflashIops, tt.wantTiflashTput)
		})
	}
}

func TestFillingCreateRequestWithDefault_ProviderCorrection(t *testing.T) {
	tests := []struct {
		name            string
		provider        globalcommonv1.CloudProvider_Name
		env             string
		wantStorageType globalngv1.ResourcePoolProfile_StorageType
		wantDataSize    uint32
		wantThroughput  uint32
	}{
		{
			name:            "alicloud clears throughput",
			provider:        globalcommonv1.CloudProvider_ALICLOUD,
			env:             "prod",
			wantStorageType: globalngv1.ResourcePoolProfile_CloudAuto,
			wantDataSize:    500,
			wantThroughput:  0,
		},
		{
			name:            "azure keeps throughput",
			provider:        globalcommonv1.CloudProvider_AZURE,
			env:             "dev",
			wantStorageType: globalngv1.ResourcePoolProfile_PremiumSSDV2,
			wantDataSize:    100,
			wantThroughput:  125,
		},
	}

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

			fillingCreateRequestWithDefault(context.Background(), request, tt.env, true)

			assertStorage(t, request.GetDemand().GetTikvData(), tt.wantStorageType, tt.wantDataSize, 3000, tt.wantThroughput)
			assertStorage(t, request.GetDemand().GetTiflashCn(), tt.wantStorageType, 500, 3000, tt.wantThroughput)
			assertStorage(t, request.GetDemand().GetTiflashWn(), tt.wantStorageType, 500, 3000, tt.wantThroughput)
		})
	}
}

func TestFillingCreateRequestWithDefault_UnknownEnvFallsBackToProviderDefault(t *testing.T) {
	request := newCreateResourcePoolRequest(globalcommonv1.CloudProvider_AWS)
	request.Demand = &globalngv1.CreateResourcePoolRequest_Demand{
		InitialRcu: 50000,
	}

	fillingCreateRequestWithDefault(context.Background(), request, "qa", true)

	assertStorage(t, request.GetDemand().GetTikvData(), globalngv1.ResourcePoolProfile_GP3, 1500, 8000, 400)
	assertStorage(t, request.GetDemand().GetTikvRaftlog(), globalngv1.ResourcePoolProfile_GP3, 40, 3000, 125)
	assertStorage(t, request.GetDemand().GetTiflashCn(), globalngv1.ResourcePoolProfile_GP3, 500, 6000, 625)
	assertStorage(t, request.GetDemand().GetTiflashWn(), globalngv1.ResourcePoolProfile_GP3, 500, 6000, 625)
}

func TestFillingCreateRequestWithDefault_NonRcuDrivenUsesProviderDefault(t *testing.T) {
	request := newCreateResourcePoolRequest(globalcommonv1.CloudProvider_AWS)
	request.Demand = &globalngv1.CreateResourcePoolRequest_Demand{
		InitialRcu: 50000,
	}

	fillingCreateRequestWithDefault(context.Background(), request, "prod", false)

	assertStorage(t, request.GetDemand().GetTikvData(), globalngv1.ResourcePoolProfile_GP3, 1500, 8000, 400)
	assertStorage(t, request.GetDemand().GetTiflashCn(), globalngv1.ResourcePoolProfile_GP3, 500, 6000, 625)
}

func TestFillingCreateRequestWithDefault_PartialDemand(t *testing.T) {
	request := newCreateResourcePoolRequest(globalcommonv1.CloudProvider_AWS)
	request.Demand = &globalngv1.CreateResourcePoolRequest_Demand{
		InitialRcu: 50000,
		TikvData: &globalngv1.CreateResourcePoolRequest_Storage{
			StorageType: globalngv1.ResourcePoolProfile_IO2,
		},
	}

	fillingCreateRequestWithDefault(context.Background(), request, "prod", true)

	assertStorage(t, request.GetDemand().GetTikvData(), globalngv1.ResourcePoolProfile_IO2, 1000, 5000, 0)
	assertStorage(t, request.GetDemand().GetTikvRaftlog(), globalngv1.ResourcePoolProfile_GP3, 40, 4000, 125)
	assertStorage(t, request.GetDemand().GetTiflashCn(), globalngv1.ResourcePoolProfile_GP3, 500, 3000, 125)
	assertStorage(t, request.GetDemand().GetTiflashWn(), globalngv1.ResourcePoolProfile_GP3, 500, 3000, 125)
}

func TestFillingCreateRequestWithDefault_DoesNotOverwriteExplicitStorageValues(t *testing.T) {
	request := newCreateResourcePoolRequest(globalcommonv1.CloudProvider_AWS)
	request.Demand = &globalngv1.CreateResourcePoolRequest_Demand{
		InitialRcu: 50000,
		TikvData: &globalngv1.CreateResourcePoolRequest_Storage{
			StorageType:        globalngv1.ResourcePoolProfile_GP3,
			DiskSizeGi:         777,
			DiskIops:           888,
			DiskThroughputMbps: 999,
		},
	}

	fillingCreateRequestWithDefault(context.Background(), request, "prod", true)

	assertStorage(t, request.GetDemand().GetTikvData(), globalngv1.ResourcePoolProfile_GP3, 777, 888, 999)
}

func TestFillingCreateRequestWithDefault_LocalSSDCorrection(t *testing.T) {
	request := newCreateResourcePoolRequest(globalcommonv1.CloudProvider_AWS)
	request.Demand = &globalngv1.CreateResourcePoolRequest_Demand{
		InitialRcu: 50000,
		TikvData: &globalngv1.CreateResourcePoolRequest_Storage{
			StorageType: globalngv1.ResourcePoolProfile_LocalSSD,
		},
	}

	fillingCreateRequestWithDefault(context.Background(), request, "prod", true)

	assertStorage(t, request.GetDemand().GetTikvData(), globalngv1.ResourcePoolProfile_LocalSSD, 0, 0, 0)
}

func TestDefaultResourcePoolInitialRcu(t *testing.T) {
	tests := []struct {
		name        string
		requiredRcu uint32
		want        uint32
	}{
		{
			name: "uses default when required rcu is lower",
			want: defaultInitialRcu,
		},
		{
			name:        "uses required rcu when higher than default",
			requiredRcu: 100000,
			want:        100000,
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			assert.Equal(t, tt.want, defaultResourcePoolInitialRcu(tt.requiredRcu))
		})
	}
}

func newCreateResourcePoolRequest(provider globalcommonv1.CloudProvider_Name) *globalngv1.CreateResourcePoolRequest {
	return &globalngv1.CreateResourcePoolRequest{
		ResourcePool: &globalngv1.ResourcePool{
			ProviderId: provider,
		},
	}
}

func assertStorage(t *testing.T,
	storage *globalngv1.CreateResourcePoolRequest_Storage,
	storageType globalngv1.ResourcePoolProfile_StorageType,
	sizeGi uint32,
	iops uint32,
	throughputMbps uint32,
) {
	t.Helper()
	require.NotNil(t, storage)
	assert.Equal(t, storageType, storage.GetStorageType())
	assert.Equal(t, sizeGi, storage.GetDiskSizeGi())
	assert.Equal(t, iops, storage.GetDiskIops())
	assert.Equal(t, throughputMbps, storage.GetDiskThroughputMbps())
}