package task

import (
	"context"
	"testing"

	"github.com/go-logr/logr/testr"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
	"github.com/tidbcloud/infra-provider/regional-server/api/v1alpha1"
	"github.com/tidbcloud/infra-provider/regional-server/pkg/config"
	"github.com/tidbcloud/infra-provider/regional-server/pkg/consts"
	"github.com/tidbcloud/infra-provider/regional-server/pkg/repository"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/runtime"
	"sigs.k8s.io/controller-runtime/pkg/client"
	"sigs.k8s.io/controller-runtime/pkg/client/fake"
)

func TestNodePool_CreateLetsDataPlaneCloudTagsOverrideGeneratedLabels(t *testing.T) {
	ctx := context.Background()
	nodePoolTask, kubeClient := newNodePoolTestTask(t)
	nodePoolTask.Config.ResourceLabels = map[string]string{
		"tags.tidbcloud.com/owner": "resource-owner",
	}
	nodePoolTask.DataPlane.Spec.CloudTags = map[string]string{
		"tags.tidbcloud.com/owner": "dataplane-owner",
		v1alpha1.RegionKey:         "custom-region",
	}

	result := nodePoolTask.Create(ctx)

	require.NoError(t, result.Error())
	require.True(t, result.IsDoneState())

	nodePool := &v1alpha1.NodePool{}
	err := kubeClient.Get(ctx, client.ObjectKey{
		Namespace: "dataplane-metadata",
		Name:      nodePoolTask.buildNodePoolName("dp-1", "m6i.2xlarge", "v1.30.0"),
	}, nodePool)
	require.NoError(t, err)
	assert.Equal(t, "dataplane-owner", nodePool.Labels["tags.tidbcloud.com/owner"])
	assert.Equal(t, "custom-region", nodePool.Labels[v1alpha1.RegionKey])
}

func TestNodePool_CreateLeavesPlacementGroupDisabledForBuiltInPools(t *testing.T) {
	// Built-in NodePools created by this controller carry
	// Spec.PlacementGroupEnabled derived from the topology predicate
	// "AWS + single-AZ?" (see docs/design/single-az-ha-placement-groups.md).
	// The default test fixture does not populate DataPlane.Spec.NodeZones,
	// so singleAzEKS returns false and the pool is created with PG=false.
	// The single-AZ / multi-AZ / non-EKS branches are exercised by the
	// dedicated tests below.
	ctx := context.Background()
	nodePoolTask, kubeClient := newNodePoolTestTask(t)
	require.NoError(t, kubeClient.Create(ctx, &v1alpha1.VPC{
		ObjectMeta: metav1.ObjectMeta{Name: "vpc-1", Namespace: "dataplane-metadata"},
		Spec:       v1alpha1.VPCSpec{Zones: []string{"us-west-2a"}},
	}))
	nodePoolTask.VpcRW = repository.NewVPCReaderWriter(kubeClient)

	result := nodePoolTask.Create(ctx)
	require.NoError(t, result.Error())
	require.True(t, result.IsDoneState())

	nodePool := &v1alpha1.NodePool{}
	require.NoError(t, kubeClient.Get(ctx, client.ObjectKey{
		Namespace: "dataplane-metadata",
		Name:      nodePoolTask.buildNodePoolName("dp-1", "m6i.2xlarge", "v1.30.0"),
	}, nodePool))
	assert.False(t, nodePool.Spec.PlacementGroupEnabled, "no NodeZones on the dataplane -> not single-AZ -> PG remains false")
}

func TestNodePool_CreateStampsPlacementGroupEnabledForSingleAzEKS(t *testing.T) {
	ctx := context.Background()
	nodePoolTask, kubeClient := newNodePoolTestTask(t)
	require.NoError(t, kubeClient.Create(ctx, &v1alpha1.VPC{
		ObjectMeta: metav1.ObjectMeta{Name: "vpc-1", Namespace: "dataplane-metadata"},
		Spec:       v1alpha1.VPCSpec{Zones: []string{"us-west-2a"}},
	}))
	nodePoolTask.VpcRW = repository.NewVPCReaderWriter(kubeClient)
	nodePoolTask.DataPlane.Spec.NodeZones = []string{"us-west-2a"}

	result := nodePoolTask.Create(ctx)
	require.NoError(t, result.Error())
	require.True(t, result.IsDoneState())

	nodePool := &v1alpha1.NodePool{}
	require.NoError(t, kubeClient.Get(ctx, client.ObjectKey{
		Namespace: "dataplane-metadata",
		Name:      nodePoolTask.buildNodePoolName("dp-1", "m6i.2xlarge", "v1.30.0"),
	}, nodePool))
	assert.True(t, nodePool.Spec.PlacementGroupEnabled, "single-AZ EKS dataplane must stamp PG=true at create time")
}

func TestNodePool_CreateLeavesPlacementGroupDisabledForMultiAzEKS(t *testing.T) {
	ctx := context.Background()
	nodePoolTask, kubeClient := newNodePoolTestTask(t)
	require.NoError(t, kubeClient.Create(ctx, &v1alpha1.VPC{
		ObjectMeta: metav1.ObjectMeta{Name: "vpc-1", Namespace: "dataplane-metadata"},
		Spec:       v1alpha1.VPCSpec{Zones: []string{"us-west-2a", "us-west-2b", "us-west-2c"}},
	}))
	nodePoolTask.VpcRW = repository.NewVPCReaderWriter(kubeClient)
	nodePoolTask.DataPlane.Spec.NodeZones = []string{"us-west-2a", "us-west-2b", "us-west-2c"}

	result := nodePoolTask.Create(ctx)
	require.NoError(t, result.Error())
	require.True(t, result.IsDoneState())

	nodePool := &v1alpha1.NodePool{}
	require.NoError(t, kubeClient.Get(ctx, client.ObjectKey{
		Namespace: "dataplane-metadata",
		Name:      nodePoolTask.buildNodePoolName("dp-1", "m6i.2xlarge", "v1.30.0"),
	}, nodePool))
	assert.False(t, nodePool.Spec.PlacementGroupEnabled, "multi-AZ EKS dataplane must not opt into partition PGs")
}

func TestNodePool_CreateLeavesPlacementGroupDisabledForACK(t *testing.T) {
	// ACK has its own single-AZ primitives; it must never opt into the
	// EKS partition-PG path even if the dataplane happens to be
	// single-AZ-shaped. ACK also drops the k8s-version segment from
	// the NodePool name (in-place upgrade reuse), so we look up the
	// single NodePool via a label selector rather than the buildName
	// helper, which assumes the EKS version-suffixed format.
	ctx := context.Background()
	nodePoolTask, kubeClient := newNodePoolTestTask(t)
	require.NoError(t, kubeClient.Create(ctx, &v1alpha1.VPC{
		ObjectMeta: metav1.ObjectMeta{Name: "vpc-1", Namespace: "dataplane-metadata"},
		Spec:       v1alpha1.VPCSpec{Zones: []string{"cn-hangzhou-a"}},
	}))
	nodePoolTask.VpcRW = repository.NewVPCReaderWriter(kubeClient)
	nodePoolTask.Config.Provider = consts.ProviderTypeACK
	nodePoolTask.DataPlane.Spec.NodeZones = []string{"cn-hangzhou-a"}

	result := nodePoolTask.Create(ctx)
	require.NoError(t, result.Error())
	require.True(t, result.IsDoneState())

	list := &v1alpha1.NodePoolList{}
	require.NoError(t, kubeClient.List(ctx, list, client.InNamespace("dataplane-metadata")))
	require.Len(t, list.Items, 1, "exactly one NodePool is created from APPBoxConfig.NodePools")
	assert.False(t, list.Items[0].Spec.PlacementGroupEnabled, "non-EKS provider must not opt into partition PGs")
}

func TestNodePool_CreateIsIdempotentWhenAllMutableFieldsMatch(t *testing.T) {
	// A reconcile on an existing built-in NodePool whose Labels,
	// Spec.Labels, Spec.Taints, and Spec.PlacementGroupEnabled already
	// match the desired state must be a no-op (no Update, resourceVersion
	// preserved) so downstream reconcilers are not re-fired on
	// idempotent retries. The pre-existing label set is populated on
	// the seed CR so nodePoolMutableEqual returns true.
	ctx := context.Background()
	nodePoolTask, kubeClient := newNodePoolTestTask(t)
	require.NoError(t, kubeClient.Create(ctx, &v1alpha1.VPC{
		ObjectMeta: metav1.ObjectMeta{Name: "vpc-1", Namespace: "dataplane-metadata"},
		Spec:       v1alpha1.VPCSpec{Zones: []string{"us-west-2a"}},
	}))
	nodePoolTask.VpcRW = repository.NewVPCReaderWriter(kubeClient)
	nodePoolTask.DataPlane.Spec.NodeZones = []string{"us-west-2a"}

	poolName := nodePoolTask.buildNodePoolName("dp-1", "m6i.2xlarge", "v1.30.0")
	require.NoError(t, kubeClient.Create(ctx, &v1alpha1.NodePool{
		ObjectMeta: metav1.ObjectMeta{
			Namespace: "dataplane-metadata",
			Name:      poolName,
			Labels: map[string]string{
				v1alpha1.RegionKey:      "us-west-2",
				v1alpha1.VPCNameKey:     "vpc-1",
				v1alpha1.ManageKey:      v1alpha1.ManageFull,
				v1alpha1.DataPlaneIDKey: "dp-1",
				v1alpha1.APPBoxNameKey:  "appbox-1",
			},
			Annotations: map[string]string{consts.AnnotationKeyBuiltinNodePool: "true"},
		},
		Spec: v1alpha1.NodePoolSpec{
			InstanceType:          "m6i.2xlarge",
			KubernetesVersion:     "v1.30.0",
			CpuCore:               8,
			AvailableCpuMilliCore: 7 * 1000,
			MemoryGi:              32,
			AvailableMemoryMi:     30 * 1024,
			NodeType:              consts.NodeTypeDedicated,
			PlacementGroupEnabled: true,
		},
	}))

	before := &v1alpha1.NodePool{}
	require.NoError(t, kubeClient.Get(ctx, client.ObjectKey{
		Namespace: "dataplane-metadata", Name: poolName,
	}, before))

	result := nodePoolTask.Create(ctx)
	require.NoError(t, result.Error())
	require.True(t, result.IsDoneState())

	after := &v1alpha1.NodePool{}
	require.NoError(t, kubeClient.Get(ctx, client.ObjectKey{
		Namespace: "dataplane-metadata", Name: poolName,
	}, after))
	assert.Equal(t, before.ResourceVersion, after.ResourceVersion,
		"idempotent reconcile must not bump resourceVersion")
}

func TestSingleAzEKS(t *testing.T) {
	cases := []struct {
		name      string
		provider  string
		nodeZones []string
		want      bool
	}{
		{"eks-single-az", consts.ProviderTypeEKS, []string{"us-east-1a"}, true},
		{"eks-multi-az", consts.ProviderTypeEKS, []string{"us-east-1a", "us-east-1b"}, false},
		{"eks-no-zones", consts.ProviderTypeEKS, nil, false},
		{"ack-single-az", consts.ProviderTypeACK, []string{"cn-hangzhou-a"}, false},
		{"aks-single-az", "aks", []string{"eastus-1"}, false},
	}
	for _, tc := range cases {
		t.Run(tc.name, func(t *testing.T) {
			assert.Equal(t, tc.want, singleAzEKS(tc.provider, tc.nodeZones))
		})
	}
}

func TestNodePool_CreateCopiesBudgetFieldsFromConfig(t *testing.T) {
	ctx := context.Background()

	t.Run("yaml without milli and max falls back to available", func(t *testing.T) {
		nodePoolTask, kubeClient := newNodePoolTestTask(t)
		result := nodePoolTask.Create(ctx)
		require.NoError(t, result.Error())
		require.True(t, result.IsDoneState())

		nodePool := &v1alpha1.NodePool{}
		require.NoError(t, kubeClient.Get(ctx, client.ObjectKey{
			Namespace: "dataplane-metadata",
			Name:      nodePoolTask.buildNodePoolName("dp-1", "m6i.2xlarge", "v1.30.0"),
		}, nodePool))
		assert.Equal(t, uint32(7000), nodePool.Spec.AvailableCpuMilliCore)
		assert.Equal(t, uint32(7000), nodePool.Spec.MaxCpuMilli)
		assert.Equal(t, uint32(30*1024), nodePool.Spec.MaxMemoryMi)
	})

	t.Run("yaml milli and max are copied onto the CR", func(t *testing.T) {
		nodePoolTask, kubeClient := newNodePoolTestTask(t)
		nodePoolTask.Config.ControllerConfig.APPBoxConfig.NodePools[0].AvailableCPUMilli = 7460
		nodePoolTask.Config.ControllerConfig.APPBoxConfig.NodePools[0].MaxCPUMilli = 7910
		nodePoolTask.Config.ControllerConfig.APPBoxConfig.NodePools[0].MaxMemoryMi = 29899

		result := nodePoolTask.Create(ctx)
		require.NoError(t, result.Error())
		require.True(t, result.IsDoneState())

		nodePool := &v1alpha1.NodePool{}
		require.NoError(t, kubeClient.Get(ctx, client.ObjectKey{
			Namespace: "dataplane-metadata",
			Name:      nodePoolTask.buildNodePoolName("dp-1", "m6i.2xlarge", "v1.30.0"),
		}, nodePool))
		assert.Equal(t, uint32(7460), nodePool.Spec.AvailableCpuMilliCore)
		assert.Equal(t, uint32(7910), nodePool.Spec.MaxCpuMilli)
		assert.Equal(t, uint32(29899), nodePool.Spec.MaxMemoryMi)
	})
}

func newNodePoolTestTask(t *testing.T) (*NodePool, client.Client) {
	t.Helper()

	scheme := runtime.NewScheme()
	require.NoError(t, v1alpha1.AddToScheme(scheme))

	appBox := &v1alpha1.APPBox{
		ObjectMeta: metav1.ObjectMeta{
			Namespace: "dataplane-metadata",
			Name:      "appbox-1",
		},
		Spec: v1alpha1.APPBoxSpec{
			KubernetesConfig: v1alpha1.KubernetesConfig{
				Version: "v1.30.0",
			},
		},
	}
	kubeClient := fake.NewClientBuilder().
		WithScheme(scheme).
		WithObjects(appBox).
		Build()
	rootConfig := config.RootConfig{
		Namespace:          "dataplane-metadata",
		ResourceNamePrefix: "tidbcloud",
		ControllerConfig: config.ControllerConfig{
			Provider: consts.ProviderTypeEKS,
			RegionConfig: config.RegionConfig{
				Name: "us-west-2",
			},
			APPBoxConfig: config.APPBoxConfig{
				NodePools: []config.NodePoolConfig{
					{
						InstanceType:      "m6i.2xlarge",
						CPU:               8,
						AvailableCPU:      7,
						MemoryGi:          32,
						AvailableMemoryMi: 30 * 1024,
					},
				},
			},
		},
	}

	taskContext := &Context{
		Scheme: scheme,
		Local:  kubeClient,
		Logger: testr.New(t),
		Config: rootConfig,
		DataPlane: &v1alpha1.DataPlane{
			ObjectMeta: metav1.ObjectMeta{
				Namespace: "dataplane-metadata",
				Name:      "dp-1",
			},
			Status: v1alpha1.DataPlaneStatus{
				Vpc: &v1alpha1.Reference{
					Namespace: "dataplane-metadata",
					Name:      "vpc-1",
				},
				AppBox: &v1alpha1.Reference{
					Namespace: "dataplane-metadata",
					Name:      "appbox-1",
				},
			},
		},
		AppBoxRW:   repository.NewAPPBoxReaderWriter(kubeClient),
		NodePoolRW: repository.NewNodePoolReaderWriter(kubeClient),
	}

	return NewNodePool(taskContext), kubeClient
}