package operation

import (
	"context"
	"testing"

	. "github.com/onsi/gomega"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/runtime"
	"sigs.k8s.io/controller-runtime/pkg/client/fake"

	"github.com/tidbcloud/cluster-controller/server/internal/util"
	cloudv1 "github.com/tidbcloud/infra-api/api/v1"
	apiv1 "github.com/tidbcloud/infra-provider/api/tidbcloud/cluster-server/v1"
	"github.com/tidbcloud/infra-provider/pkg/consts"
	utiltime "github.com/tidbcloud/infra-provider/pkg/util/time"
)

func awsTestCluster() *cloudv1.Cluster {
	return &cloudv1.Cluster{
		ObjectMeta: metav1.ObjectMeta{Name: "c1", Namespace: "ns1"},
		Spec: cloudv1.ClusterSpec{
			Kubernetes: cloudv1.Kubernetes{Provider: cloudv1.ProviderTypeEKS},
		},
	}
}

func TestBuildStartZonalShiftWF(t *testing.T) {
	g := NewWithT(t)
	s := &Server{}
	wf := &cloudv1.Workflow{}
	err := s.buildStartZonalShiftWF(wf, &apiv1.Operation_StartZonalShift{
		AvailabilityZone: "us-east-1a",
		ExpiresIn:        "10m",
	}, awsTestCluster())
	g.Expect(err).NotTo(HaveOccurred())
	g.Expect(wf.Spec.Steps).To(HaveLen(2))
	g.Expect(wf.Spec.Steps[0].Uses).To(Equal(consts.WfActionStartZonalShiftAction))
	g.Expect(wf.Spec.Steps[0].Inputs.Object["availabilityZone"]).To(Equal("us-east-1a"))
	g.Expect(wf.Spec.Steps[1].Uses).To(Equal(consts.WorkflowActionMutateAnnotations))
	g.Expect(wf.Spec.Steps[1].Inputs.Object["type"]).To(Equal("cluster"))
	g.Expect(wf.Spec.Steps[1].Inputs.Object["annotations"]).To(Equal(map[string]interface{}{
		consts.AnnoKeyAZFaultActive: "us-east-1a",
	}))
}

func TestCreateOperation_StopZonalShift_RejectsMismatchedArgs(t *testing.T) {
	g := NewWithT(t)
	s := runtime.NewScheme()
	g.Expect(cloudv1.AddToScheme(s)).To(Succeed())

	cluster := util.BuildTestClusterCR("c1", "ns1", func(c *cloudv1.Cluster) {
		c.Spec.Kubernetes.Provider = cloudv1.ProviderTypeEKS
	})
	kubeClient := fake.NewClientBuilder().WithScheme(s).WithObjects(cluster).Build()

	_, err := NewServer(kubeClient, utiltime.RealClock{}).CreateOperation(context.TODO(), &apiv1.CreateOperationRequest{
		ClusterId: "c1",
		Operation: &apiv1.Operation{
			IdempotencyKey: "stop-shift",
			Type:           apiv1.Operation_STOP_ZONAL_SHIFT,
			ClusterId:      "c1",
			Args: &apiv1.Operation_OfflineAz{
				OfflineAz: &apiv1.Operation_OfflineAZ{AvailabilityZone: "us-east-1a"},
			},
		},
	})
	g.Expect(err).To(HaveOccurred())
	g.Expect(err.Error()).To(ContainSubstring("stop_zonal_shift args are required"))
}

func TestBuildStopZonalShiftWF_DefaultWait(t *testing.T) {
	g := NewWithT(t)
	s := &Server{}
	wf := &cloudv1.Workflow{}
	err := s.buildStopZonalShiftWF(wf, &apiv1.Operation_StopZonalShift{}, awsTestCluster())
	g.Expect(err).NotTo(HaveOccurred())
	g.Expect(wf.Spec.Steps).To(HaveLen(4))
	g.Expect(wf.Spec.Steps[0].Uses).To(Equal(consts.WfActionCancelZonalShiftAction))
	g.Expect(wf.Spec.Steps[1].Uses).To(Equal(consts.WorkflowActionSleep))
	g.Expect(wf.Spec.Steps[1].Inputs.Object["duration"]).To(Equal("30m"))
	g.Expect(wf.Spec.Steps[2].Uses).To(Equal(consts.WfActionCleanupZonalShiftAction))
	g.Expect(wf.Spec.Steps[3].Uses).To(Equal(consts.WorkflowActionMutateAnnotations))
	g.Expect(wf.Spec.Steps[3].Inputs.Object["annotations"]).To(Equal(map[string]interface{}{
		consts.AnnoKeyAZFaultActive: "-",
	}))
}

func TestBuildUpdateZonalShiftWF(t *testing.T) {
	g := NewWithT(t)
	s := &Server{}
	wf := &cloudv1.Workflow{}
	err := s.buildUpdateZonalShiftWF(wf, &apiv1.Operation_UpdateZonalShift{
		ExpiresIn: "30m",
	}, awsTestCluster())
	g.Expect(err).NotTo(HaveOccurred())
	g.Expect(wf.Spec.Steps).To(HaveLen(1))
	g.Expect(wf.Spec.Steps[0].Uses).To(Equal(consts.WfActionUpdateZonalShiftAction))
	g.Expect(wf.Spec.Steps[0].Inputs.Object["expiresIn"]).To(Equal("30m"))
	g.Expect(wf.Spec.Steps[0].Inputs.Object["zonalShiftId"]).To(BeNil())
}

func TestBuildStartZonalShiftWF_RejectsNonAWS(t *testing.T) {
	g := NewWithT(t)
	s := &Server{}
	wf := &cloudv1.Workflow{}
	cluster := awsTestCluster()
	cluster.Spec.Kubernetes.Provider = cloudv1.ProviderTypeGKE
	err := s.buildStartZonalShiftWF(wf, &apiv1.Operation_StartZonalShift{
		AvailabilityZone: "us-east-1a",
		ExpiresIn:        "10m",
	}, cluster)
	g.Expect(err).To(HaveOccurred())
}

func TestBuildStartZonalShiftWF_WithGroupIDs(t *testing.T) {
	g := NewWithT(t)
	s := &Server{}
	wf := &cloudv1.Workflow{}
	cluster := awsTestCluster()
	cluster.Annotations = map[string]string{
		cloudv1.DefaultLoadBalancerIDAnnotationKey: "default-gid",
	}
	cluster.Spec.ExtraGroups = []cloudv1.ExtraGroupSpec{{Name: "analytics", TiDB: cloudv1.TiDB{Replicas: 1}}}
	err := s.buildStartZonalShiftWF(wf, &apiv1.Operation_StartZonalShift{
		AvailabilityZone: "us-east-1a",
		ExpiresIn:        "10m",
		GroupIds:         []string{"analytics"},
	}, cluster)
	g.Expect(err).NotTo(HaveOccurred())
	g.Expect(wf.Spec.Steps[0].Inputs.Object["groupIds"]).To(Equal([]interface{}{"analytics"}))
	// InputsWrapper DeepCopyJSON must not panic on groupIds.
	g.Expect(func() { _ = wf.Spec.Steps[0].Inputs.DeepCopy() }).NotTo(Panic())
}

func TestValidateGroupIDs_Unknown(t *testing.T) {
	g := NewWithT(t)
	cluster := awsTestCluster()
	cluster.Annotations = map[string]string{
		cloudv1.DefaultLoadBalancerIDAnnotationKey: "default-gid",
	}
	err := validateGroupIDs(cluster, []string{"nope"})
	g.Expect(err).To(HaveOccurred())
}