package reconciler

import (
	"context"
	"testing"
	"time"

	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
	"github.com/tidbcloud/aws-shared-provider/pkg/tidbworker/resourcemanager/poolmanager/factory"
	"github.com/tidbcloud/aws-shared-provider/pkg/tidbworker/resourcemanager/testutils"
	rm "github.com/tidbcloud/aws-shared-provider/pkg/tidbworker/resourcemanager/types"
	v1 "k8s.io/api/core/v1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"sigs.k8s.io/controller-runtime/pkg/client"
)

func mustNewReconciler(re *require.Assertions) (rm.Reconciler, testutils.MockManager) {
	mockConfig := testutils.DefaultConfig()
	mockManager, err := testutils.NewMockManager(mockConfig)
	re.NoError(err)
	scaler := &v1.Pod{
		TypeMeta: metav1.TypeMeta{
			Kind:       "Pod",
			APIVersion: "v1",
		},
		ObjectMeta: metav1.ObjectMeta{
			Name:      "test-pod",
			Namespace: mockConfig.Namespace,
			Labels: map[string]string{
				"control-plane": "scaler",
			},
		},
		Spec: v1.PodSpec{
			Containers: []v1.Container{
				{
					Name:  "test-container",
					Image: "test-image",
				},
			},
		},
	}
	re.NoError(mockManager.GetClusterClient().Create(context.Background(), scaler))

	outsidePod := &v1.Pod{
		TypeMeta: metav1.TypeMeta{
			Kind:       "Pod",
			APIVersion: "v1",
		},
		ObjectMeta: metav1.ObjectMeta{
			Name:      "outside-pod",
			Namespace: "outside",
			Labels: map[string]string{
				"control-plane": "outside",
			},
		},
		Spec: v1.PodSpec{
			Containers: []v1.Container{
				{
					Name:  "test-container",
					Image: "test-image",
				},
			},
		},
	}
	re.NoError(mockManager.GetClusterClient().Create(context.Background(), outsidePod))

	return NewReconciler(mockManager, &rm.ReconcilerConfig{
		ReconcileInterval: 300 * time.Millisecond,
	}), mockManager
}

func TestReconciler(t *testing.T) {
	re := require.New(t)
	reconciler, manager := mustNewReconciler(re)

	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()
	go func() {
		re.NoError(reconciler.Start(ctx))
	}()

	// Create worker and services from the config.
	re.EventuallyWithT(func(c *assert.CollectT) {
		mustMatchWorkers(c, manager)
		mustPreserveScaler(c, manager)
		mustPreserveOutside(c, manager)
	}, 3*time.Second, 500*time.Millisecond)

	//  Remove a pool from pools config, reconciler should remove pods and services from the cluster.
	cfg := manager.GetConfig()
	for pool := range cfg.PoolsConfig.Pools {
		delete(cfg.PoolsConfig.Pools, pool)
		break
	}
	re.NoError(manager.UpdateCluster(context.Background(), cfg))
	re.EventuallyWithT(func(c *assert.CollectT) {
		mustMatchWorkers(c, manager)
		mustPreserveScaler(c, manager)
		mustPreserveOutside(c, manager)
	}, 3*time.Second, 500*time.Millisecond)

	// Remove excessive workers and services when shrinking pool size.
	for pool := range cfg.PoolsConfig.Pools {
		cfg.PoolsConfig.Pools[pool].Size--
		break
	}
	re.NoError(manager.UpdateCluster(context.Background(), cfg))
	re.EventuallyWithT(func(c *assert.CollectT) {
		mustMatchWorkers(c, manager)
		mustPreserveScaler(c, manager)
		mustPreserveOutside(c, manager)
	}, 3*time.Second, 500*time.Millisecond)

	// Add new workers and services when expanding pool size.
	for pool := range cfg.PoolsConfig.Pools {
		cfg.PoolsConfig.Pools[pool].Size++
		break
	}
	re.NoError(manager.UpdateCluster(context.Background(), cfg))
	re.EventuallyWithT(func(c *assert.CollectT) {
		mustMatchWorkers(c, manager)
		mustPreserveScaler(c, manager)
		mustPreserveOutside(c, manager)
	}, 3*time.Second, 500*time.Millisecond)

	// Recreate missing workers after manual deletion.
	cli := manager.GetClusterClient()
	re.NoError(cli.DeleteAllOf(ctx, &v1.Pod{},
		client.InNamespace(manager.GetNamespace()),
		client.MatchingLabels(factory.WorkerLabels()),
	))
	re.EventuallyWithT(func(c *assert.CollectT) {
		mustMatchWorkers(c, manager)
		mustPreserveScaler(c, manager)
		mustPreserveOutside(c, manager)
	}, 3*time.Second, 500*time.Millisecond)
}

// mustMatchWorkers asserts that the pod and service objects in the cluster matches those stated in the config.
func mustMatchWorkers(c *assert.CollectT, manager testutils.MockManager) {
	mustMatchServices(c, manager)
	mustMatchPods(c, manager)
}

// mustMatchServices asserts that the service in the cluster matches those stated in the config.
func mustMatchServices(c *assert.CollectT, manager testutils.MockManager) {
	expectedServices := manager.GetExpectedServices()
	cli := manager.GetClusterClient()
	namespace := manager.GetNamespace()

	// Check service spec.
	var totalExpected int
	for _, poolServices := range expectedServices {
		for _, expected := range poolServices {
			totalExpected++
			actual := &v1.Service{}
			err := cli.Get(context.Background(), client.ObjectKey{
				Namespace: namespace,
				Name:      expected.Name,
			}, actual)
			assert.NoError(c, err)
			assert.Equal(c, expected.Spec, actual.Spec)
		}
	}

	// Check service count.
	allServices := &v1.ServiceList{}
	assert.NoError(c, cli.List(context.Background(), allServices, client.InNamespace(namespace)))
	assert.Len(c, allServices.Items, totalExpected)
}

// mustMatchPods asserts that the pod in the cluster matches those stated in the config.
func mustMatchPods(c *assert.CollectT, manager testutils.MockManager) {
	expectedPods, err := manager.GetExpectedPods()
	assert.NoError(c, err)
	cli := manager.GetClusterClient()
	namespace := manager.GetNamespace()

	// Check pod spec.
	totalExpected := 1 // The scaler pod.
	for _, poolPods := range expectedPods {
		for _, expected := range poolPods {
			totalExpected++
			actual := &v1.Pod{}
			err := cli.Get(context.Background(), client.ObjectKey{
				Namespace: namespace,
				Name:      expected.Name,
			}, actual)
			assert.NoError(c, err)
			assert.Equal(c, expected.Spec, actual.Spec)
		}
	}
	// Check pod count.
	allPods := &v1.PodList{}
	assert.NoError(c, cli.List(context.Background(), allPods, client.InNamespace(namespace)))
	assert.Len(c, allPods.Items, totalExpected)
}

func mustPreserveScaler(c *assert.CollectT, manager testutils.MockManager) {
	cli := manager.GetClusterClient()
	namespace := manager.GetNamespace()
	scaler := &v1.Pod{}
	err := cli.Get(context.Background(), client.ObjectKey{
		Namespace: namespace,
		Name:      "test-pod",
	}, scaler)
	assert.NoError(c, err)
}

func mustPreserveOutside(c *assert.CollectT, manager testutils.MockManager) {
	cli := manager.GetClusterClient()
	outside := &v1.Pod{}
	err := cli.Get(context.Background(), client.ObjectKey{
		Namespace: "outside",
		Name:      "outside-pod",
	}, outside)
	assert.NoError(c, err)
}