package br

import (
	"context"
	"testing"

	"github.com/stretchr/testify/mock"
	"github.com/stretchr/testify/require"
	dplanemgrv1 "github.com/tidbcloud/cluster-service-ng/api/dataplane-manager/v1"
	dataplanemgrmocks "github.com/tidbcloud/cluster-service-ng/pkg/dataplanemgr/mocks"
	schemav1 "github.com/tidbcloud/cluster-service-ng/pkg/storage/etcd/schema/v1"
)

func TestGetTiBRNodePlacementSelectsSmallestFittingActivePool(t *testing.T) {
	client := dataplanemgrmocks.NewClient(t)
	client.EXPECT().
		GetDataPlane(mock.Anything, &dplanemgrv1.GetDataPlaneRequest{DataPlaneId: "dp-1"}).
		Return(&dplanemgrv1.DataPlane{
			Tags:       map[string]string{schemav1.LabelUserTenantKey: "tenant-1"},
			Kubernetes: &dplanemgrv1.DataPlane_Kubernetes{Version: "1.33"},
		}, nil).
		Once()
	client.EXPECT().
		ListNodePools(mock.Anything, &dplanemgrv1.ListNodePoolsRequest{DataPlaneId: "dp-1"}).
		Return(&dplanemgrv1.ListNodePoolsResponse{NodePools: []*dplanemgrv1.NodePool{
			newTiBRNodePool("large", "1.33", 8, 16, 7900, 15000),
			newTiBRNodePool("small", "1.33", 2, 8, 1900, 7000),
			newTiBRNodePool("too-small", "1.33", 1, 4, 900, 3500),
			func() *dplanemgrv1.NodePool {
				pool := newTiBRNodePool("shared", "1.33", 2, 8, 1900, 7000)
				pool.NodeType = dplanemgrv1.NodeType_Shared
				return pool
			}(),
			func() *dplanemgrv1.NodePool {
				pool := newTiBRNodePool("local-disk", "1.33", 2, 8, 1900, 7000)
				pool.WithLocalDisk = true
				return pool
			}(),
			func() *dplanemgrv1.NodePool {
				pool := newTiBRNodePool("inactive", "1.33", 2, 8, 1900, 7000)
				pool.Status = dplanemgrv1.NodePoolStatus_NODEPOOL_UPDATING
				return pool
			}(),
			newTiBRNodePool("old-version", "1.32", 2, 8, 1900, 7000),
		}}, nil).
		Once()

	tenant, placement, err := getTiBRNodePlacement(context.Background(), client, "dp-1", "1", "4812Mi")
	require.NoError(t, err)
	require.Equal(t, "tenant-1", tenant)
	require.Equal(t, TiBRNodePlacement{NodeSpec: "2C8G", NodeVersion: "1.33"}, placement)
}

func TestResolveTiBRNodePlacementRejectsMissingFit(t *testing.T) {
	client := dataplanemgrmocks.NewClient(t)
	client.EXPECT().
		ListNodePools(mock.Anything, &dplanemgrv1.ListNodePoolsRequest{DataPlaneId: "dp-1"}).
		Return(&dplanemgrv1.ListNodePoolsResponse{NodePools: []*dplanemgrv1.NodePool{
			newTiBRNodePool("small", "1.33", 2, 8, 1000, 4096),
		}}, nil).
		Once()

	_, err := resolveTiBRNodePlacement(context.Background(), client, "dp-1", "1.33", "2", "8Gi")
	require.ErrorContains(t, err, "no active dedicated node pool fits TiBR")
}

func TestResolveTiBRNodePlacementUsesMemoryAsTieBreaker(t *testing.T) {
	client := dataplanemgrmocks.NewClient(t)
	client.EXPECT().
		ListNodePools(mock.Anything, mock.Anything).
		Return(&dplanemgrv1.ListNodePoolsResponse{NodePools: []*dplanemgrv1.NodePool{
			newTiBRNodePool("memory-large", "1.33", 4, 32, 3900, 31000),
			newTiBRNodePool("memory-small", "1.33", 4, 16, 3900, 15000),
		}}, nil).
		Once()

	placement, err := resolveTiBRNodePlacement(context.Background(), client, "dp-1", "1.33", "3", "10Gi")
	require.NoError(t, err)
	require.Equal(t, "4C16G", placement.NodeSpec)
}

func newTiBRNodePool(name, version string, cpu, memory, availableCPU, availableMemory uint32) *dplanemgrv1.NodePool {
	return &dplanemgrv1.NodePool{
		Name:                  name,
		K8SVersion:            version,
		CpuCore:               cpu,
		MemoryGi:              memory,
		AvailableCpuMilliCore: availableCPU,
		AvailableMemoryMi:     availableMemory,
		NodeType:              dplanemgrv1.NodeType_Dedicated,
		Status:                dplanemgrv1.NodePoolStatus_NODEPOOL_ACTIVE,
	}
}