package imports

import (
	"context"
	"sync"
	"testing"

	"github.com/go-logr/logr"
	"github.com/stretchr/testify/require"
	managerconfig "github.com/tidbcloud/aws-shared-provider/cmd/manager/app/config"
	providerscheme "github.com/tidbcloud/aws-shared-provider/pkg/scheme"
	"github.com/tidbcloud/aws-shared-provider/pkg/utils/labels"
	cloudv1 "github.com/tidbcloud/infra-api/api/v1"
	infra "github.com/tidbcloud/tidb-management-service/api/serverless_infra"
	corev1 "k8s.io/api/core/v1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/types"
	"sigs.k8s.io/controller-runtime/pkg/client"
	"sigs.k8s.io/controller-runtime/pkg/client/fake"
)

type importCMCountingClient struct {
	client.Client

	mu           sync.Mutex
	importCMGets int
}

func (c *importCMCountingClient) Get(
	ctx context.Context,
	key client.ObjectKey,
	obj client.Object,
	opts ...client.GetOption,
) error {
	if key.Name == importCmName {
		if _, ok := obj.(*corev1.ConfigMap); ok {
			c.mu.Lock()
			c.importCMGets++
			c.mu.Unlock()
		}
	}
	return c.Client.Get(ctx, key, obj, opts...)
}

func (c *importCMCountingClient) getCount() int {
	c.mu.Lock()
	defer c.mu.Unlock()
	return c.importCMGets
}

type contextBlockingClient struct {
	client.Client
}

func (c *contextBlockingClient) Get(
	ctx context.Context,
	_ client.ObjectKey,
	_ client.Object,
	_ ...client.GetOption,
) error {
	<-ctx.Done()
	return ctx.Err()
}

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

	tests := []struct {
		name             string
		tidbPool         string
		annotations      map[string]string
		config           map[string]string
		wantSelector     string
		wantDedicatedCPU string
		wantToleration   string
		withoutSelector  bool
		withoutTaint     bool
	}{
		{
			name:           "default scheduling",
			wantSelector:   "import-pool",
			wantToleration: "import-pool",
		},
		{
			name:           "vip default scheduling",
			tidbPool:       vipTidbPool,
			wantSelector:   "import-pool",
			wantToleration: "import-pool",
		},
		{
			name:             "super vip default scheduling",
			tidbPool:         superVipTidbPool,
			wantSelector:     "dedicated",
			wantDedicatedCPU: "64",
			wantToleration:   "dedicated",
		},
		{
			name:     "super vip dedicated scheduling is configurable",
			tidbPool: superVipTidbPool,
			config: map[string]string{
				superVipNodeSelectorKey:                 "dedicated",
				superVipTolerationsKey:                  "dedicated",
				labels.ServerlessDedicatedNodeCpuNumKey: "64",
			},
			wantSelector:     "dedicated",
			wantDedicatedCPU: "64",
			wantToleration:   "dedicated",
		},
		{
			name:         "only node selector",
			config:       map[string]string{nodeSelectorKey: "custom-worker"},
			wantSelector: "custom-worker",
			withoutTaint: true,
		},
		{
			name:            "only toleration",
			config:          map[string]string{tolerationsKey: "custom-taint"},
			wantToleration:  "custom-taint",
			withoutSelector: true,
		},
		{
			name:           "node selector and toleration",
			config:         map[string]string{nodeSelectorKey: "custom-worker", tolerationsKey: "custom-taint"},
			wantSelector:   "custom-worker",
			wantToleration: "custom-taint",
		},
		{
			name:           "vip configuration",
			tidbPool:       canaryVipTidbPool,
			config:         map[string]string{vipNodeSelectorKey: "vip-worker", vipTolerationsKey: "vip-taint"},
			wantSelector:   "vip-worker",
			wantToleration: "vip-taint",
		},
		{
			name:         "super vip annotation selects super vip configuration",
			annotations:  map[string]string{importSuperVIPAnno: "true"},
			config:       map[string]string{superVipNodeSelectorKey: "large-worker"},
			wantSelector: "large-worker",
			withoutTaint: true,
		},
		{
			name:           "empty values use legacy defaults",
			config:         map[string]string{nodeSelectorKey: "", tolerationsKey: ""},
			wantSelector:   "import-pool",
			wantToleration: "import-pool",
		},
	}

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

			env := &ServerlessTierEnv{
				importObj: &cloudv1.Import{
					ObjectMeta: metav1.ObjectMeta{Annotations: tt.annotations},
				},
				clusterObj: &infra.Cluster{},
			}
			env.clusterObj.Status.TidbPool = tt.tidbPool
			podSpec := &corev1.PodSpec{}

			applyImportScheduling(env, podSpec, &corev1.ConfigMap{Data: tt.config})

			if tt.withoutSelector {
				require.Empty(t, podSpec.NodeSelector)
			} else {
				wantNodeSelector := map[string]string{
					labels.ServerlessComponentNodeLabelKey: tt.wantSelector,
				}
				if tt.wantDedicatedCPU != "" {
					wantNodeSelector[labels.ServerlessDedicatedNodeCpuNumKey] = tt.wantDedicatedCPU
				}
				require.Equal(t, wantNodeSelector, podSpec.NodeSelector)
			}
			if tt.withoutTaint {
				require.Empty(t, podSpec.Tolerations)
			} else {
				require.Equal(t, []corev1.Toleration{{
					Key:      labels.ServerlessComponentNodeLabelKey,
					Operator: corev1.TolerationOpEqual,
					Effect:   corev1.TaintEffectNoSchedule,
					Value:    tt.wantToleration,
				}}, podSpec.Tolerations)
			}
		})
	}
}

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

	tests := []struct {
		name string
		key  string
	}{
		{name: "invalid cpu request", key: cpuResourceRequestKey},
		{name: "invalid cpu limit", key: cpuResourceLimitKey},
		{name: "invalid memory request", key: memoryResourceRequestKey},
		{name: "invalid memory limit", key: memoryResourceLimitKey},
	}

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

			env := &ServerlessTierEnv{
				importObj:  &cloudv1.Import{},
				clusterObj: &infra.Cluster{},
			}

			_, err := getLightningResources(env, &corev1.ConfigMap{
				Data: map[string]string{tt.key: "not-a-resource-quantity"},
			})
			require.Error(t, err)
			require.ErrorContains(t, err, tt.key)
		})
	}
}

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

	ctx := context.Background()
	managerCfg := &managerconfig.Config{
		CloudProviderType:        "aws-free-tier",
		ImportJobsNamespace:      "import-jobs",
		ImportSAName:             "serverless-import",
		ServerlessLightningImage: "lightning",
	}
	importCm := &corev1.ConfigMap{
		ObjectMeta: metav1.ObjectMeta{Name: importCmName, Namespace: managerCfg.ImportJobsNamespace},
		Data:       map[string]string{enableCheckpointKey: "true"},
	}
	regionalClient := &importCMCountingClient{
		Client: fake.NewClientBuilder().
			WithScheme(providerscheme.Scheme).
			WithObjects(importCm).
			Build(),
	}
	importObj := &cloudv1.Import{
		ObjectMeta: metav1.ObjectMeta{Name: "import-1", Namespace: "project"},
		Spec: cloudv1.ImportSpec{
			LightningConfig: "[lightning]\ntable-concurrency = 1\n",
		},
	}
	infraClient := fake.NewClientBuilder().
		WithScheme(providerscheme.Scheme).
		WithStatusSubresource(&cloudv1.Import{}).
		WithObjects(importObj).
		Build()
	reconciledImport := &cloudv1.Import{}
	require.NoError(t, infraClient.Get(ctx, types.NamespacedName{
		Name: importObj.Name, Namespace: importObj.Namespace,
	}, reconciledImport))

	clusterObj := &infra.Cluster{}
	clusterObj.Status.KeyspaceToken = "keyspace-1"
	service := &ImportService{
		infraCli: infraClient,
		env: &ServerlessTierEnv{
			importObj:      reconciledImport,
			clusterObj:     clusterObj,
			regionalEksCli: regionalClient,
			managerCfg:     managerCfg,
		},
		logger: logr.Discard(),
	}

	require.NoError(t, service.Create(ctx, reconciledImport))
	require.Equal(t, 1, regionalClient.getCount())
}

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

	ctx, cancel := context.WithCancel(context.Background())
	cancel()
	env := &ServerlessTierEnv{
		regionalEksCli: &contextBlockingClient{},
		managerCfg:     &managerconfig.Config{ImportJobsNamespace: "import-jobs"},
	}

	_, err := env.getImportCm(ctx)
	require.ErrorIs(t, err, context.Canceled)
}