package resourcepoolcapacity

import (
	"testing"

	"github.com/stretchr/testify/require"
	clusterngv1beta "github.com/tidbcloud/cluster-service-ng/api/cluster-service-ng/v1beta"
	"google.golang.org/protobuf/proto"

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

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

	tests := []struct {
		name       string
		components []*clusterngv1beta.Component
		assert     func(*testing.T, *clusterngv1beta.SharedPool)
	}{
		{
			name: "projects TiKV and ignores hidden components",
			components: []*clusterngv1beta.Component{
				component(
					clusterngv1beta.ComponentType_TIKV,
					8,
					32,
					3,
					storage(clusterngv1beta.Component_Storage_DATA, 500),
					storage(clusterngv1beta.Component_Storage_RAFTLOG, 40),
				),
				{
					ComponentType: clusterngv1beta.ComponentType_PD,
				},
				component(
					clusterngv1beta.ComponentType_TIFLASH_WRITE,
					16,
					64,
					2,
					storage(clusterngv1beta.Component_Storage_DATA, 999),
				),
			},
			assert: func(t *testing.T, sharedPool *clusterngv1beta.SharedPool) {
				t.Helper()

				capacity, err := Project(sharedPool.GetComponentGroups())
				require.NoError(t, err)
				require.EqualValues(t, 8, capacity.GetTikv().GetNodeSpec().GetCpu())
				require.EqualValues(t, 32, capacity.GetTikv().GetNodeSpec().GetMemoryGi())
				require.EqualValues(t, 3, capacity.GetTikv().GetNodeCount())
				require.Len(t, capacity.GetTikv().GetStorages(), 1)
				require.Equal(t, "data", capacity.GetTikv().GetStorages()[0].GetName())
				require.EqualValues(t, 500, capacity.GetTikv().GetStorages()[0].GetSizeGi())
				require.Nil(t, capacity.GetTiflash())
			},
		},
		{
			name: "omits missing TiFlash Compute",
			components: []*clusterngv1beta.Component{
				component(clusterngv1beta.ComponentType_TIKV, 8, 32, 3),
			},
			assert: func(t *testing.T, sharedPool *clusterngv1beta.SharedPool) {
				t.Helper()

				capacity, err := Project(sharedPool.GetComponentGroups())
				require.NoError(t, err)
				require.Nil(t, capacity.GetTiflash())
			},
		},
		{
			name: "omits zero-count TiFlash Compute without requiring a snapshot",
			components: []*clusterngv1beta.Component{
				component(clusterngv1beta.ComponentType_TIKV, 8, 32, 3),
				{ComponentType: clusterngv1beta.ComponentType_TIFLASH_COMPUTE},
			},
			assert: func(t *testing.T, sharedPool *clusterngv1beta.SharedPool) {
				t.Helper()

				capacity, err := Project(sharedPool.GetComponentGroups())
				require.NoError(t, err)
				require.Nil(t, capacity.GetTiflash())
			},
		},
		{
			name: "projects positive TiFlash Compute",
			components: []*clusterngv1beta.Component{
				component(clusterngv1beta.ComponentType_TIKV, 8, 32, 3),
				component(
					clusterngv1beta.ComponentType_TIFLASH_COMPUTE,
					16,
					64,
					1,
					storage(clusterngv1beta.Component_Storage_DATA, 500),
				),
			},
			assert: func(t *testing.T, sharedPool *clusterngv1beta.SharedPool) {
				t.Helper()

				capacity, err := Project(sharedPool.GetComponentGroups())
				require.NoError(t, err)
				require.EqualValues(t, 16, capacity.GetTiflash().GetNodeSpec().GetCpu())
				require.EqualValues(t, 64, capacity.GetTiflash().GetNodeSpec().GetMemoryGi())
				require.EqualValues(t, 1, capacity.GetTiflash().GetNodeCount())
				require.Equal(t, "data", capacity.GetTiflash().GetStorages()[0].GetName())
				require.EqualValues(t, 500, capacity.GetTiflash().GetStorages()[0].GetSizeGi())
			},
		},
	}

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

			test.assert(t, &clusterngv1beta.SharedPool{ComponentGroups: test.components})
		})
	}
}

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

	tests := []struct {
		name        string
		components  []*clusterngv1beta.Component
		errContains string
	}{
		{
			name:        "nil component",
			components:  []*clusterngv1beta.Component{nil, component(clusterngv1beta.ComponentType_TIKV, 8, 32, 3)},
			errContains: "component is nil",
		},
		{
			name:        "missing TiKV",
			components:  []*clusterngv1beta.Component{},
			errContains: "exactly one TIKV component is required",
		},
		{
			name: "duplicate TiKV",
			components: []*clusterngv1beta.Component{
				component(clusterngv1beta.ComponentType_TIKV, 8, 32, 3),
				component(clusterngv1beta.ComponentType_TIKV, 16, 64, 3),
			},
			errContains: "duplicate TIKV component",
		},
		{
			name: "duplicate TiFlash Compute",
			components: []*clusterngv1beta.Component{
				component(clusterngv1beta.ComponentType_TIKV, 8, 32, 3),
				component(clusterngv1beta.ComponentType_TIFLASH_COMPUTE, 16, 64, 1),
				component(clusterngv1beta.ComponentType_TIFLASH_COMPUTE, 16, 64, 1),
			},
			errContains: "duplicate TIFLASH_COMPUTE component",
		},
		{
			name: "TiKV nil node quota",
			components: []*clusterngv1beta.Component{
				{ComponentType: clusterngv1beta.ComponentType_TIKV},
			},
			errContains: "TIKV node quota is required",
		},
		{
			name: "TiKV nil node spec",
			components: []*clusterngv1beta.Component{
				{
					ComponentType: clusterngv1beta.ComponentType_TIKV,
					NodeQuota:     &clusterngv1beta.NodeQuota{NodeCount: 3},
				},
			},
			errContains: "TIKV node spec is required",
		},
		{
			name: "TiKV zero CPU",
			components: []*clusterngv1beta.Component{
				component(clusterngv1beta.ComponentType_TIKV, 0, 32, 3),
			},
			errContains: "TIKV node spec CPU and memory must be positive",
		},
		{
			name: "TiKV zero memory",
			components: []*clusterngv1beta.Component{
				component(clusterngv1beta.ComponentType_TIKV, 8, 0, 3),
			},
			errContains: "TIKV node spec CPU and memory must be positive",
		},
		{
			name: "TiKV zero node count",
			components: []*clusterngv1beta.Component{
				component(clusterngv1beta.ComponentType_TIKV, 8, 32, 0),
			},
			errContains: "TIKV node count must be positive",
		},
		{
			name: "nil storage",
			components: []*clusterngv1beta.Component{
				component(clusterngv1beta.ComponentType_TIKV, 8, 32, 3, nil),
			},
			errContains: "TIKV storage is nil",
		},
		{
			name: "unknown TiKV storage",
			components: []*clusterngv1beta.Component{
				component(
					clusterngv1beta.ComponentType_TIKV,
					8,
					32,
					3,
					storage(clusterngv1beta.Component_Storage_SLOWLOG, 20),
				),
			},
			errContains: "unsupported TIKV storage SLOWLOG",
		},
		{
			name: "unknown TiFlash Compute storage",
			components: []*clusterngv1beta.Component{
				component(clusterngv1beta.ComponentType_TIKV, 8, 32, 3),
				component(
					clusterngv1beta.ComponentType_TIFLASH_COMPUTE,
					16,
					64,
					1,
					storage(clusterngv1beta.Component_Storage_RAFTLOG, 40),
				),
			},
			errContains: "unsupported TIFLASH_COMPUTE storage RAFTLOG",
		},
		{
			name: "duplicate storage",
			components: []*clusterngv1beta.Component{
				component(
					clusterngv1beta.ComponentType_TIKV,
					8,
					32,
					3,
					storage(clusterngv1beta.Component_Storage_DATA, 500),
					storage(clusterngv1beta.Component_Storage_DATA, 600),
				),
			},
			errContains: "duplicate TIKV storage DATA",
		},
	}

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

			capacity, err := Project(test.components)
			require.Nil(t, capacity)
			require.ErrorContains(t, err, test.errContains)
		})
	}
}

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

	tikv := normalizedScaleTarget(4, 16, 3,
		normalizedStorage("data", 500),
		normalizedStorage("raftlog", 40),
	)
	zeroCompute := normalizedScaleTarget(8, 32, 0, normalizedStorage("data", 500))

	omitted, err := ProjectScaleTargets(tikv, zeroCompute, false)
	require.NoError(t, err)
	require.Nil(t, omitted.GetTiflash())
	require.Len(t, omitted.GetTikv().GetStorages(), 1)
	require.Equal(t, "data", omitted.GetTikv().GetStorages()[0].GetName())

	explicitZero, err := ProjectScaleTargets(tikv, zeroCompute, true)
	require.NoError(t, err)
	require.NotNil(t, explicitZero.GetTiflash())
	require.Zero(t, explicitZero.GetTiflash().GetNodeCount())
	require.Equal(t, "data", explicitZero.GetTiflash().GetStorages()[0].GetName())
}

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

	tikv := normalizedScaleTarget(4, 16, 3,
		normalizedStorage("DATA", 500),
		normalizedStorage("data", 600),
	)
	capacity, err := ProjectScaleTargets(tikv, nil, false)
	require.Nil(t, capacity)
	require.ErrorContains(t, err, "duplicate TIKV storage DATA")

	tikv = normalizedScaleTarget(4, 16, 3, normalizedStorage("SLOWLOG", 20))
	capacity, err = ProjectScaleTargets(tikv, nil, false)
	require.Nil(t, capacity)
	require.ErrorContains(t, err, "unsupported TIKV storage SLOWLOG")

	tikv = normalizedScaleTarget(4, 16, 3, &globalngv1.ByocSharedComponentStorage{StorageName: "DATA"})
	capacity, err = ProjectScaleTargets(tikv, nil, false)
	require.Nil(t, capacity)
	require.ErrorContains(t, err, "normalized storage is incomplete")
}

func normalizedScaleTarget(
	cpu uint32,
	memoryGi uint32,
	nodeCount uint32,
	storages ...*globalngv1.ByocSharedComponentStorage,
) *globalngv1.ByocSharedComponentScaleTarget {
	return &globalngv1.ByocSharedComponentScaleTarget{
		Cpu:       proto.Uint32(cpu),
		MemoryGi:  proto.Uint32(memoryGi),
		NodeCount: proto.Uint32(nodeCount),
		Storages:  storages,
	}
}

func normalizedStorage(name string, sizeGi uint32) *globalngv1.ByocSharedComponentStorage {
	return &globalngv1.ByocSharedComponentStorage{
		StorageName: name,
		SizeGi:      proto.Uint32(sizeGi),
	}
}

func component(
	componentType clusterngv1beta.ComponentType,
	cpu uint32,
	memoryGi uint32,
	nodeCount uint32,
	storages ...*clusterngv1beta.Component_Storage,
) *clusterngv1beta.Component {
	return &clusterngv1beta.Component{
		ComponentType: componentType,
		NodeQuota: &clusterngv1beta.NodeQuota{
			NodeSpec: &clusterngv1beta.NodeSpec{
				Cpu:      cpu,
				MemoryGi: memoryGi,
			},
			NodeCount: nodeCount,
		},
		Storages: storages,
	}
}

func storage(name clusterngv1beta.Component_Storage_Name, sizeGi uint32) *clusterngv1beta.Component_Storage {
	return &clusterngv1beta.Component_Storage{
		Name:   name,
		SizeGi: sizeGi,
	}
}