package common

import (
	"context"
	"errors"
	"testing"

	opv1alpha1 "github.com/pingcap/tidb-operator/pkg/apis/pingcap/v1alpha1"
	"github.com/stretchr/testify/require"
	appsv1 "k8s.io/api/apps/v1"
	apierrors "k8s.io/apimachinery/pkg/api/errors"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/runtime/schema"
	"k8s.io/apimachinery/pkg/types"
	"sigs.k8s.io/controller-runtime/pkg/client"
	"sigs.k8s.io/controller-runtime/pkg/client/fake"
	"sigs.k8s.io/controller-runtime/pkg/controller/controllerutil"

	"github.com/tidbcloud/infra-provider/pkg/scheme"
)

func TestCreateOrUpdateTidbClusterWithRetryPreservesConcurrentUpdates(t *testing.T) {
	ctx := context.Background()
	existing := &opv1alpha1.TidbCluster{
		ObjectMeta: metav1.ObjectMeta{
			Namespace: "ns",
			Name:      "db",
			Annotations: map[string]string{
				"existing": "keep",
			},
		},
		Spec: opv1alpha1.TidbClusterSpec{
			TiDB: &opv1alpha1.TiDBSpec{Replicas: 1},
		},
	}
	baseClient := fake.NewClientBuilder().WithScheme(scheme.Scheme).WithObjects(existing).Build()
	cli := &conflictOnceTidbClusterClient{Client: baseClient}

	obj := &opv1alpha1.TidbCluster{
		ObjectMeta: metav1.ObjectMeta{
			Namespace: "ns",
			Name:      "db",
		},
	}
	mutateCalls := 0
	opResult, err := createOrUpdateTidbClusterWithRetry(ctx, cli, obj, func() error {
		mutateCalls++
		if obj.Annotations == nil {
			obj.Annotations = map[string]string{}
		}
		obj.Annotations["cluster-controller"] = "desired"
		if obj.Spec.TiDB == nil {
			obj.Spec.TiDB = &opv1alpha1.TiDBSpec{}
		}
		obj.Spec.TiDB.Replicas = 2
		return nil
	})
	require.NoError(t, err)
	require.Equal(t, controllerutil.OperationResultUpdated, opResult)
	require.Equal(t, 2, mutateCalls)
	require.Equal(t, 2, cli.updateCalls)

	got := &opv1alpha1.TidbCluster{}
	require.NoError(t, cli.Get(ctx, types.NamespacedName{Namespace: "ns", Name: "db"}, got))
	require.Equal(t, int32(2), got.Spec.TiDB.Replicas)
	require.Equal(t, "desired", got.Annotations["cluster-controller"])
	require.Equal(t, "keep", got.Annotations["existing"])
	require.Equal(t, "kept", got.Annotations["tidb-operator"])
	require.NotNil(t, got.Status.PD.StatefulSet)
	require.Equal(t, int32(3), got.Status.PD.StatefulSet.Replicas)
}

type conflictOnceTidbClusterClient struct {
	client.Client
	updateCalls int
}

func (c *conflictOnceTidbClusterClient) Update(ctx context.Context, obj client.Object, opts ...client.UpdateOption) error {
	c.updateCalls++
	if c.updateCalls == 1 {
		latest := &opv1alpha1.TidbCluster{}
		if err := c.Client.Get(ctx, client.ObjectKeyFromObject(obj), latest); err != nil {
			return err
		}
		if latest.Annotations == nil {
			latest.Annotations = map[string]string{}
		}
		latest.Annotations["tidb-operator"] = "kept"
		latest.Status.PD.SetStatefulSet(&appsv1.StatefulSetStatus{Replicas: 3})
		if err := c.Client.Update(ctx, latest, opts...); err != nil {
			return err
		}
		return apierrors.NewConflict(schema.GroupResource{Group: "pingcap.com", Resource: "tidbclusters"}, obj.GetName(), errors.New("injected conflict"))
	}
	return c.Client.Update(ctx, obj, opts...)
}