package service

import (
	"context"
	"testing"

	opv2 "github.com/pingcap/tidb-operator/api/v2/core/v1alpha1"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
	apiv1 "github.com/tidbcloud/cluster-service-ng/api/tidbcloud/cluster-agent/v1"
	corev1 "k8s.io/api/core/v1"
	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"
)

type stubNodeRestarter struct {
	providerID string
}

func (s *stubNodeRestarter) Restart(_ context.Context, providerID string) error {
	s.providerID = providerID
	return nil
}

func TestAgentService_RestartNodeAcceptsTiFlashWritePodLabel(t *testing.T) {
	scheme := runtime.NewScheme()
	require.NoError(t, corev1.AddToScheme(scheme))

	const (
		clusterID   = "2056254137925898240"
		podName     = "db-93739ff3-write-tiflash-367apm"
		nodeName    = "ip-10-0-0-1"
		providerID  = "aws:///us-west-1a/i-1234567890abcdef0"
		operationID = "restartnode001"
	)

	pod := &corev1.Pod{
		ObjectMeta: metav1.ObjectMeta{
			Name:      podName,
			Namespace: buildNamespace(clusterID),
			Labels: map[string]string{
				K8sComponentLabelKey: "write-tiflash",
				K8sInstanceLabelKey:  "db-93739ff3-write",
			},
		},
		Spec: corev1.PodSpec{
			NodeName: nodeName,
		},
	}
	node := &corev1.Node{
		ObjectMeta: metav1.ObjectMeta{
			Name: nodeName,
		},
		Spec: corev1.NodeSpec{
			ProviderID: providerID,
		},
	}

	restarter := &stubNodeRestarter{}
	cli := fake.NewClientBuilder().WithScheme(scheme).WithObjects(pod, node).Build()
	service := NewAgentService(cli, nil, WithNodeRestarter(restarter))

	resp, err := service.restartNode(context.Background(), &apiv1.RestartReq{
		ClusterId:   clusterID,
		TargetName:  podName,
		TargetType:  apiv1.RestartTargetType_NODE,
		OperationId: operationID,
	})

	require.NoError(t, err)
	require.NotNil(t, resp)
	assert.Equal(t, "success", resp.Message)
	assert.Equal(t, providerID, restarter.providerID)

	var updatedNode corev1.Node
	require.NoError(t, cli.Get(context.Background(), client.ObjectKey{Name: nodeName}, &updatedNode))
	assert.Equal(t, operationID, updatedNode.Annotations[operationIDAnnotationKey])
}

func TestAgentService_RestartTiProxyGroupPreservesMaxSurge(t *testing.T) {
	scheme := runtime.NewScheme()
	require.NoError(t, opv2.AddToScheme(scheme))

	const (
		clusterID   = "2056254137925898240"
		groupName   = "db-93739ff3"
		operationID = "restarttiproxy001"
	)
	maxSurge := int32(1024)

	group := &opv2.TiProxyGroup{
		ObjectMeta: metav1.ObjectMeta{
			Name:      groupName,
			Namespace: buildNamespace(clusterID),
		},
		Spec: opv2.TiProxyGroupSpec{
			MaxSurge: &maxSurge,
		},
		Status: opv2.TiProxyGroupStatus{
			CommonStatus: opv2.CommonStatus{
				Conditions: []metav1.Condition{
					{Type: string(opv2.CondReady), Status: metav1.ConditionTrue},
					{Type: string(opv2.CondSynced), Status: metav1.ConditionTrue},
				},
			},
		},
	}

	cli := fake.NewClientBuilder().
		WithScheme(scheme).
		WithObjects(group).
		Build()
	service := NewAgentService(cli, nil)

	resp, err := service.restartComponent(context.Background(), &apiv1.RestartReq{
		ClusterId:   clusterID,
		TargetType:  apiv1.RestartTargetType_COMPONENT,
		TargetName:  "tiproxy",
		TargetGroup: groupName,
		OperationId: operationID,
	})

	require.NoError(t, err)
	require.NotNil(t, resp)
	assert.Equal(t, "success", resp.Message)

	var updated opv2.TiProxyGroup
	require.NoError(t, cli.Get(context.Background(), client.ObjectKey{Namespace: buildNamespace(clusterID), Name: groupName}, &updated))
	assert.Equal(t, operationID, updated.Spec.Template.Annotations[operationIDAnnotationKey])
	require.NotNil(t, updated.Spec.MaxSurge)
	assert.Equal(t, maxSurge, *updated.Spec.MaxSurge)
}