package gkeheterogeneous

import (
	"context"
	"testing"

	"github.com/stretchr/testify/require"
	apierrors "k8s.io/apimachinery/pkg/api/errors"
	"k8s.io/apimachinery/pkg/apis/meta/v1/unstructured"
	"sigs.k8s.io/controller-runtime/pkg/client"
	"sigs.k8s.io/controller-runtime/pkg/client/fake"

	cloudv1 "github.com/tidbcloud/infra-api/api/v1"
	"github.com/tidbcloud/infra-provider/pkg/consts"
)

func testGatekeeperAssignRef(name, labelComputeClassName, specComputeClassName string) *unstructured.Unstructured {
	obj := &unstructured.Unstructured{
		Object: map[string]interface{}{
			"apiVersion": gatekeeperMutationGroup + "/" + gatekeeperMutationVersion,
			"kind":       gatekeeperAssignKind,
			"metadata": map[string]interface{}{
				"name": name,
			},
			"spec": map[string]interface{}{
				"parameters": map[string]interface{}{
					"assign": map[string]interface{}{
						"value": map[string]interface{}{},
					},
				},
			},
		},
	}
	if labelComputeClassName != "" {
		obj.SetLabels(map[string]string{GatekeeperAssignComputeClassLabelKey: labelComputeClassName})
	}
	if specComputeClassName != "" {
		_ = unstructured.SetNestedField(obj.Object, specComputeClassName, "spec", "parameters", "assign", "value", consts.GKEComputeClassLabelKey)
	}
	return obj
}

func TestGatekeeperAssignRefsForComputeClass(t *testing.T) {
	target := testGatekeeperAssignRef("target", "old-priority", "")
	specOnly := testGatekeeperAssignRef("spec-only", "", "target-priority")
	other := testGatekeeperAssignRef("other", "other-priority", "other-priority")
	cli := fake.NewClientBuilder().WithRuntimeObjects(target, specOnly, other).Build()

	refs, err := GatekeeperAssignRefsForComputeClass(context.Background(), cli, "target-priority")
	require.NoError(t, err)
	require.Len(t, refs, 1)
	require.Equal(t, "spec-only", refs[0].Name)
	require.Empty(t, refs[0].LabelComputeClassName)
	require.Equal(t, "target-priority", refs[0].SpecComputeClassName)

	refs, err = GatekeeperAssignRefsForComputeClass(context.Background(), cli, "old-priority")
	require.NoError(t, err)
	require.Len(t, refs, 1)
	require.Equal(t, "target", refs[0].Name)
	require.Equal(t, "old-priority", refs[0].LabelComputeClassName)
}

func TestGatekeeperAssignRefsForComputeClassFailsClosedForManagedMalformedSpec(t *testing.T) {
	malformed := testGatekeeperAssignRef("malformed", "target-priority", "")
	require.NoError(t, unstructured.SetNestedField(malformed.Object, []interface{}{"target-priority"}, "spec", "parameters", "assign", "value", consts.GKEComputeClassLabelKey))
	cli := fake.NewClientBuilder().WithRuntimeObjects(malformed).Build()

	_, err := GatekeeperAssignRefsForComputeClass(context.Background(), cli, "target-priority")
	require.Error(t, err)
	require.Contains(t, err.Error(), "malformed")
}

func TestGatekeeperAssignRefsForComputeClassSkipsUnmanagedMalformedSpec(t *testing.T) {
	malformed := testGatekeeperAssignRef("malformed", "", "")
	require.NoError(t, unstructured.SetNestedField(malformed.Object, []interface{}{"target-priority"}, "spec", "parameters", "assign", "value", consts.GKEComputeClassLabelKey))
	target := testGatekeeperAssignRef("target", "", "target-priority")
	cli := fake.NewClientBuilder().WithRuntimeObjects(malformed, target).Build()

	refs, err := GatekeeperAssignRefsForComputeClass(context.Background(), cli, "target-priority")
	require.NoError(t, err)
	require.Len(t, refs, 1)
	require.Equal(t, "target", refs[0].Name)
}

func TestGatekeeperAssignRefsForComputeClassSkipsOtherLabelMalformedSpec(t *testing.T) {
	malformed := testGatekeeperAssignRef("malformed", "other-priority", "")
	require.NoError(t, unstructured.SetNestedField(malformed.Object, []interface{}{"other-priority"}, "spec", "parameters", "assign", "value", consts.GKEComputeClassLabelKey))
	target := testGatekeeperAssignRef("target", "", "target-priority")
	cli := fake.NewClientBuilder().WithRuntimeObjects(malformed, target).Build()

	refs, err := GatekeeperAssignRefsForComputeClass(context.Background(), cli, "target-priority")
	require.NoError(t, err)
	require.Len(t, refs, 1)
	require.Equal(t, "target", refs[0].Name)
}

func TestGatekeeperAssignRefsForComputeClassSkipsTopologySpreadCompanion(t *testing.T) {
	companion := testGatekeeperAssignRef("topology-spread", "", "")
	companion.SetLabels(map[string]string{
		consts.ManagedByLabelKey:        ManagedByClusterAgentLabelValue,
		HeterogeneousDeployLabelKey:     GatekeeperAssignDeployLabelValue,
		GatekeeperAssignPurposeLabelKey: GatekeeperAssignPurposeTopologySpread,
	})
	require.NoError(t, unstructured.SetNestedField(companion.Object, gatekeeperTopologySpreadNodeAffinityLocation, "spec", "location"))
	require.NoError(t, unstructured.SetNestedField(companion.Object, "Ignore", "spec", "parameters", "assign", "value"))
	target := testGatekeeperAssignRef("target", "", "target-priority")
	cli := fake.NewClientBuilder().WithRuntimeObjects(companion, target).Build()

	refs, err := GatekeeperAssignRefsForComputeClass(context.Background(), cli, "target-priority")
	require.NoError(t, err)
	require.Len(t, refs, 1)
	require.Equal(t, "target", refs[0].Name)
}

func TestGatekeeperAssignRefsForComputeClassDoesNotSkipTopologySpreadWithComputeClassLabel(t *testing.T) {
	assign := testGatekeeperAssignRef("topology-spread-with-label", "target-priority", "")
	assign.SetLabels(map[string]string{
		consts.ManagedByLabelKey:             ManagedByClusterAgentLabelValue,
		HeterogeneousDeployLabelKey:          GatekeeperAssignDeployLabelValue,
		GatekeeperAssignPurposeLabelKey:      GatekeeperAssignPurposeTopologySpread,
		GatekeeperAssignComputeClassLabelKey: "target-priority",
	})
	require.NoError(t, unstructured.SetNestedField(assign.Object, gatekeeperTopologySpreadNodeAffinityLocation, "spec", "location"))
	cli := fake.NewClientBuilder().WithRuntimeObjects(assign).Build()

	refs, err := GatekeeperAssignRefsForComputeClass(context.Background(), cli, "target-priority")
	require.NoError(t, err)
	require.Len(t, refs, 1)
	require.Equal(t, "topology-spread-with-label", refs[0].Name)
	require.Equal(t, "target-priority", refs[0].LabelComputeClassName)
}

func TestGatekeeperAssignRefsForComputeClassDoesNotSkipTopologySpreadWithComputeClassSpec(t *testing.T) {
	assign := testGatekeeperAssignRef("topology-spread-with-spec", "", "target-priority")
	assign.SetLabels(map[string]string{
		consts.ManagedByLabelKey:        ManagedByClusterAgentLabelValue,
		HeterogeneousDeployLabelKey:     GatekeeperAssignDeployLabelValue,
		GatekeeperAssignPurposeLabelKey: GatekeeperAssignPurposeTopologySpread,
	})
	require.NoError(t, unstructured.SetNestedField(assign.Object, gatekeeperTopologySpreadNodeAffinityLocation, "spec", "location"))
	cli := fake.NewClientBuilder().WithRuntimeObjects(assign).Build()

	refs, err := GatekeeperAssignRefsForComputeClass(context.Background(), cli, "target-priority")
	require.NoError(t, err)
	require.Len(t, refs, 1)
	require.Equal(t, "topology-spread-with-spec", refs[0].Name)
	require.Equal(t, "target-priority", refs[0].SpecComputeClassName)
}

func testManagedGatekeeperAssign(t *testing.T, clusterID string, componentType cloudv1.ComponentType, tidbGroupID, suffix, purpose string) *unstructured.Unstructured {
	t.Helper()
	obj := &unstructured.Unstructured{
		Object: map[string]interface{}{
			"apiVersion": gatekeeperMutationGroup + "/" + gatekeeperMutationVersion,
			"kind":       gatekeeperAssignKind,
			"metadata": map[string]interface{}{
				"name": ResourceName(clusterID, componentType, tidbGroupID, suffix),
			},
			"spec": map[string]interface{}{
				"match": map[string]interface{}{
					"scope":      "Namespaced",
					"namespaces": []interface{}{tidbClusterNamespacePrefix + clusterID},
				},
			},
		},
	}
	labels := gatekeeperAssignComponentLabels(clusterID, componentType, tidbGroupID)
	labels[GatekeeperAssignPurposeLabelKey] = purpose
	obj.SetLabels(labels)
	if purpose == GatekeeperAssignPurposeTopologySpread {
		require.NoError(t, unstructured.SetNestedField(obj.Object, gatekeeperTopologySpreadNodeAffinityLocation, "spec", "location"))
		require.NoError(t, unstructured.SetNestedField(obj.Object, gatekeeperTopologySpreadNodeAffinityIgnore, "spec", "parameters", "assign", "value"))
	}
	return obj
}

func TestDeleteGatekeeperAssignForComponentIfExistsDeletesPrimaryAndCompanion(t *testing.T) {
	primary := testManagedGatekeeperAssign(t, "123", cloudv1.ComponentTiKV, "", gatekeeperAssignPrimarySuffix, GatekeeperAssignPurposeNodeSelector)
	companion := testManagedGatekeeperAssign(t, "123", cloudv1.ComponentTiKV, "", gatekeeperAssignTopologySpreadSuffix, GatekeeperAssignPurposeTopologySpread)
	cli := fake.NewClientBuilder().WithRuntimeObjects(primary, companion).Build()

	err := DeleteGatekeeperAssignForComponentIfExists(context.Background(), cli, "123", cloudv1.ComponentTiKV, "")
	require.NoError(t, err)

	gotPrimary := newGatekeeperAssign()
	err = cli.Get(context.Background(), clientObjectKey(primary.GetName()), gotPrimary)
	require.True(t, apierrors.IsNotFound(err))
	gotCompanion := newGatekeeperAssign()
	err = cli.Get(context.Background(), clientObjectKey(companion.GetName()), gotCompanion)
	require.True(t, apierrors.IsNotFound(err))
}

func TestDeleteGatekeeperAssignForComponentIfExistsMissingPrimaryIsNoop(t *testing.T) {
	cli := fake.NewClientBuilder().Build()

	err := DeleteGatekeeperAssignForComponentIfExists(context.Background(), cli, "123", cloudv1.ComponentTiKV, "")
	require.NoError(t, err)
}

func TestDeleteGatekeeperAssignForComponentIfExistsRejectsEmptyTiDBGroupID(t *testing.T) {
	group1 := testManagedGatekeeperAssign(t, "123", cloudv1.ComponentTiDB, "group1", gatekeeperAssignPrimarySuffix, GatekeeperAssignPurposeNodeSelector)
	cli := fake.NewClientBuilder().WithRuntimeObjects(group1).Build()

	err := DeleteGatekeeperAssignForComponentIfExists(context.Background(), cli, "123", cloudv1.ComponentTiDB, "")
	require.Error(t, err)
	require.Contains(t, err.Error(), "tidb group id is required")

	gotGroup1 := newGatekeeperAssign()
	require.NoError(t, cli.Get(context.Background(), clientObjectKey(group1.GetName()), gotGroup1))
}

func TestDeleteGatekeeperAssignForComponentIfExistsDeletesOrphanCompanion(t *testing.T) {
	companion := testManagedGatekeeperAssign(t, "123", cloudv1.ComponentTiKV, "", gatekeeperAssignTopologySpreadSuffix, GatekeeperAssignPurposeTopologySpread)
	cli := fake.NewClientBuilder().WithRuntimeObjects(companion).Build()

	err := DeleteGatekeeperAssignForComponentIfExists(context.Background(), cli, "123", cloudv1.ComponentTiKV, "")
	require.NoError(t, err)

	gotCompanion := newGatekeeperAssign()
	err = cli.Get(context.Background(), clientObjectKey(companion.GetName()), gotCompanion)
	require.True(t, apierrors.IsNotFound(err))
}

func TestDeleteGatekeeperAssignForComponentIfExistsRejectsOwnerMismatch(t *testing.T) {
	primary := testManagedGatekeeperAssign(t, "123", cloudv1.ComponentTiKV, "", gatekeeperAssignPrimarySuffix, GatekeeperAssignPurposeNodeSelector)
	labels := primary.GetLabels()
	labels[consts.ClusterLabelKey] = "456"
	primary.SetLabels(labels)
	cli := fake.NewClientBuilder().WithRuntimeObjects(primary).Build()

	err := DeleteGatekeeperAssignForComponentIfExists(context.Background(), cli, "123", cloudv1.ComponentTiKV, "")
	require.Error(t, err)
	require.Contains(t, err.Error(), consts.ClusterLabelKey)
}

func TestDeleteGatekeeperAssignForComponentIfExistsRejectsMultiplePrimaries(t *testing.T) {
	primary := testManagedGatekeeperAssign(t, "123", cloudv1.ComponentTiKV, "", gatekeeperAssignPrimarySuffix, GatekeeperAssignPurposeNodeSelector)
	duplicate := testManagedGatekeeperAssign(t, "123", cloudv1.ComponentTiKV, "", "duplicate-compute-class-assign", GatekeeperAssignPurposeNodeSelector)
	cli := fake.NewClientBuilder().WithRuntimeObjects(primary, duplicate).Build()

	err := DeleteGatekeeperAssignForComponentIfExists(context.Background(), cli, "123", cloudv1.ComponentTiKV, "")
	require.Error(t, err)
	require.Contains(t, err.Error(), "multiple primary")
}

func TestDeleteGatekeeperAssignForComponentIfExistsDeletesOnlyRequestedTiDBGroup(t *testing.T) {
	group1 := testManagedGatekeeperAssign(t, "123", cloudv1.ComponentTiDB, "group1", gatekeeperAssignPrimarySuffix, GatekeeperAssignPurposeNodeSelector)
	group2 := testManagedGatekeeperAssign(t, "123", cloudv1.ComponentTiDB, "group2", gatekeeperAssignPrimarySuffix, GatekeeperAssignPurposeNodeSelector)
	cli := fake.NewClientBuilder().WithRuntimeObjects(group1, group2).Build()

	err := DeleteGatekeeperAssignForComponentIfExists(context.Background(), cli, "123", cloudv1.ComponentTiDB, "group1")
	require.NoError(t, err)

	gotGroup1 := newGatekeeperAssign()
	err = cli.Get(context.Background(), clientObjectKey(group1.GetName()), gotGroup1)
	require.True(t, apierrors.IsNotFound(err))
	gotGroup2 := newGatekeeperAssign()
	require.NoError(t, cli.Get(context.Background(), clientObjectKey(group2.GetName()), gotGroup2))
}

func clientObjectKey(name string) client.ObjectKey {
	return client.ObjectKey{Name: name}
}