package audit

import (
	"context"
	"strconv"
	"testing"

	v1alpha1 "github.com/pingcap/tidb-operator/pkg/apis/pingcap/v1alpha1"
	"github.com/stretchr/testify/require"
	apps "k8s.io/api/apps/v1"
	corev1 "k8s.io/api/core/v1"
	storagev1 "k8s.io/api/storage/v1"
	apierrors "k8s.io/apimachinery/pkg/api/errors"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/runtime"
	"k8s.io/apimachinery/pkg/types"
	ctrl "sigs.k8s.io/controller-runtime"
	"sigs.k8s.io/controller-runtime/pkg/client"
	"sigs.k8s.io/controller-runtime/pkg/client/fake"

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

var scheme *runtime.Scheme

func init() {
	scheme = runtime.NewScheme()
	_ = v1alpha1.AddToScheme(scheme)
	_ = cloudv1.AddToScheme(scheme)
	_ = corev1.AddToScheme(scheme)
	_ = apps.AddToScheme(scheme)
	_ = storagev1.AddToScheme(scheme)
}

func TestEnsureAuditLogStorageClass(t *testing.T) {
	audit := &cloudv1.Audit{
		Enable: true,
	}
	tcClient := fake.NewClientBuilder().WithScheme(scheme).Build()
	key := types.NamespacedName{Name: auditLogStorageClassName}

	awsNet := &cloudv1.Network{
		Spec: cloudv1.NetworkSpec{
			Provider: "aws",
		},
	}
	gcpNet := &cloudv1.Network{
		Spec: cloudv1.NetworkSpec{
			Provider: "gcp",
		},
	}
	unknownNet := &cloudv1.Network{
		Spec: cloudv1.NetworkSpec{
			Provider: "unknown",
		},
	}

	t.Run("s3", func(t *testing.T) {
		audit.BucketURL = "s3://yinzong/niubi"
		err := ensureAuditLogStorageClass(context.Background(), tcClient, awsNet)
		require.Nil(t, err)
		sc := new(storagev1.StorageClass)
		err = tcClient.Get(context.Background(), key, sc)
		require.Nil(t, err)
		require.Equal(t, "gp3", sc.Parameters["type"])
		require.Equal(t, corev1.PersistentVolumeReclaimDelete, *sc.ReclaimPolicy)
		err = tcClient.Delete(context.Background(), sc)
		require.Nil(t, err)
	})

	t.Run("gcs", func(t *testing.T) {
		audit.BucketURL = "gcs://yinzong/niubi"
		err := ensureAuditLogStorageClass(context.Background(), tcClient, gcpNet)
		require.Nil(t, err)
		sc := new(storagev1.StorageClass)
		err = tcClient.Get(context.Background(), key, sc)
		require.Nil(t, err)
		require.Equal(t, "pd-ssd", sc.Parameters["type"])
		require.Equal(t, corev1.PersistentVolumeReclaimDelete, *sc.ReclaimPolicy)
		err = tcClient.Delete(context.Background(), sc)
		require.Nil(t, err)
	})

	t.Run("unknown", func(t *testing.T) {
		audit.BucketURL = "unknown://yinzong/niubi"
		err := ensureAuditLogStorageClass(context.Background(), tcClient, unknownNet)
		require.NotNil(t, err)
	})
}

func TestReconcile(t *testing.T) {
	ctx := context.Background()
	var tcClient client.Client
	var infraClient client.Client
	var network *cloudv1.Network
	var cluster *cloudv1.Cluster
	var tc *v1alpha1.TidbCluster

	reset := func(t *testing.T) {
		tcClient = fake.NewClientBuilder().WithScheme(scheme).Build()
		infraClient = fake.NewClientBuilder().WithScheme(scheme).Build()

		network = &cloudv1.Network{
			Spec: cloudv1.NetworkSpec{
				Provider: "aws",
			},
			Status: cloudv1.NetworkStatus{
				IAM: cloudv1.IAMStatus{
					ExternalId: "eid",
				},
			},
		}

		cluster = &cloudv1.Cluster{
			ObjectMeta: metav1.ObjectMeta{
				Namespace: "default",
				Name:      "c",
			},
			Spec: cloudv1.ClusterSpec{
				Cluster: cloudv1.ClusterTopology{
					TiDB: cloudv1.TiDB{
						Replicas: 3,
					},
				},
			},
		}

		tc = &v1alpha1.TidbCluster{
			ObjectMeta: metav1.ObjectMeta{
				Namespace: "default",
				Name:      "tc",
			},
			Spec: v1alpha1.TidbClusterSpec{
				TiDB: &v1alpha1.TiDBSpec{
					Replicas: 3,
				},
			},
		}

		err := infraClient.Create(ctx, cluster)
		require.Nil(t, err)
		// fake client do not set this so we set it ourself.
		cluster.CreationTimestamp = metav1.Now()
		require.False(t, cluster.CreationTimestamp.IsZero())
		err = tcClient.Create(ctx, tc)
		require.Nil(t, err)
		tc.CreationTimestamp = metav1.Now()
		require.False(t, tc.CreationTimestamp.IsZero())
	}

	t.Run("test no op", func(t *testing.T) {
		reset(t)

		res, err := Reconcile(ctx, network, cluster, infraClient, tc, tcClient, false)
		require.Equal(t, ctrl.Result{}, res)
		require.Nil(t, err)
	})

	t.Run("test enable", func(t *testing.T) {
		reset(t)
		tk := toolKit{cli: tcClient, T: t}

		cluster.Spec.Cluster.TiDB.Audit = &cloudv1.Audit{
			Enable:    true,
			BucketURL: "s3://yinzong/niubi",
			Region:    "us-west2",
			RoleArn:   "aws::sdf:dfdf",
		}

		err := infraClient.Update(ctx, cluster)
		require.Nil(t, err)

		// update tc spec to enable
		res, err := Reconcile(ctx, network, cluster, infraClient, tc, tcClient, false)
		require.Nil(t, err)
		require.NotEmpty(t, res)
		tcGet := new(v1alpha1.TidbCluster)
		err = tcClient.Get(ctx, types.NamespacedName{Namespace: tc.Namespace, Name: tc.Name}, tcGet)
		require.Nil(t, err)
		checkAuditEnable(t, tcGet)

		// old sts will be deleted
		tk.createTiDBSTS(tc, false)
		res, err = Reconcile(ctx, network, cluster, infraClient, tc, tcClient, false)
		require.Nil(t, err)
		require.NotEmpty(t, res)
		stsKey := types.NamespacedName{Namespace: tc.Namespace, Name: tc.Name + "-tidb"}
		sts := new(apps.StatefulSet)
		err = tcClient.Get(ctx, stsKey, sts)
		require.NotNil(t, err)
		require.True(t, apierrors.IsNotFound(err))
		// should not delete new sts
		tk.createTiDBSTS(tc, true)
		res, err = Reconcile(ctx, network, cluster, infraClient, tc, tcClient, false)
		require.Nil(t, err)
		require.NotEmpty(t, res)
		err = tcClient.Get(ctx, stsKey, sts)
		require.Nil(t, err)

		// check delete pod one by one
		for i := 0; i < int(tc.Spec.TiDB.Replicas); i++ {
			tk.replacePod(tc, i, true /*ready*/, false /*withVolume*/)
		}
		// trigger delete pod-0
		res, err = Reconcile(ctx, network, cluster, infraClient, tc, tcClient, false)
		require.Nil(t, err)
		require.NotEmpty(t, res)
		require.Nil(t, tk.getPod(tc, 0))
		// do not skip to next pod(pod-1) if previous one not ready(pod-0).
		tk.replacePod(tc, 0, false, true)
		require.Nil(t, err)
		require.NotEmpty(t, res)
		require.NotNil(t, tk.getPod(tc, 1))
		// pod-0 ready and delete pod-1
		tk.replacePod(tc, 0, true, true)
		res, err = Reconcile(ctx, network, cluster, infraClient, tc, tcClient, false)
		require.Nil(t, err)
		require.NotEmpty(t, res)
		require.Nil(t, tk.getPod(tc, 1))
		// pod-1 ready and delete pod-2
		tk.replacePod(tc, 1, true, true)
		res, err = Reconcile(ctx, network, cluster, infraClient, tc, tcClient, false)
		require.Nil(t, err)
		require.NotEmpty(t, res)
		require.Nil(t, tk.getPod(tc, 2))
		// pod-2 ready, all pods updated
		tk.replacePod(tc, 2, true, true)
		res, err = Reconcile(ctx, network, cluster, infraClient, tc, tcClient, false)
		require.Nil(t, err)
		require.Empty(t, res)
		require.True(t, getAuditReady(cluster))
	})

	t.Run("test disable", func(t *testing.T) {
		reset(t)

		// enable and driver to finish reconcile
		cluster.Spec.Cluster.TiDB.Audit = &cloudv1.Audit{
			Enable:    true,
			BucketURL: "s3://yinzong/niubi",
			Region:    "us-west2",
			RoleArn:   "aws::sdf:dfdf",
		}
		err := infraClient.Update(ctx, cluster)
		require.Nil(t, err)
		_, err = Reconcile(ctx, network, cluster, infraClient, tc, tcClient, false)
		require.Nil(t, err)
		tcGet := new(v1alpha1.TidbCluster)
		err = tcClient.Get(ctx, types.NamespacedName{Namespace: tc.Namespace, Name: tc.Name}, tcGet)
		require.Nil(t, err)
		checkAuditEnable(t, tcGet)

		cluster.Spec.Cluster.TiDB.Audit.Enable = false
		err = infraClient.Update(ctx, cluster)
		require.Nil(t, err)
		res, err := Reconcile(ctx, network, cluster, infraClient, tc, tcClient, false)
		require.Nil(t, err)
		require.Empty(t, res)
		tcGet = new(v1alpha1.TidbCluster)
		err = tcClient.Get(ctx, types.NamespacedName{Namespace: tc.Namespace, Name: tc.Name}, tcGet)
		require.Nil(t, err)
		checkAuditDisable(t, tc)
		require.False(t, getAuditReady(cluster))
	})
}

func checkAuditEnable(t *testing.T, tc *v1alpha1.TidbCluster) {
	t.Helper()

	require.Equal(t, slowlogVolumeName, tc.Spec.TiDB.SlowLogVolumeName)
	require.True(t, containStorageVolume(tc.Spec.TiDB.StorageVolumes, slowlogVolumeName))
	require.True(t, containContainer(tc.Spec.TiDB.AdditionalContainers, auditLogName))
	require.True(t, containString(tc.Spec.TiDB.Plugins, pluginName))
}

func checkAuditDisable(t *testing.T, tc *v1alpha1.TidbCluster) {
	t.Helper()

	require.Equal(t, "", tc.Spec.TiDB.SlowLogVolumeName)
	require.False(t, containStorageVolume(tc.Spec.TiDB.StorageVolumes, auditLogName))
	require.False(t, containContainer(tc.Spec.TiDB.AdditionalContainers, auditLogName))
	require.False(t, containString(tc.Spec.TiDB.Plugins, pluginName))
}

type toolKit struct {
	cli client.Client
	*testing.T
}

func (tk *toolKit) createTiDBSTS(tc *v1alpha1.TidbCluster, withVolume bool) {
	sts := &apps.StatefulSet{
		ObjectMeta: metav1.ObjectMeta{
			Namespace: tc.Namespace,
			Name:      tc.Name + "-tidb",
		},
	}

	if withVolume {
		sts.Spec.VolumeClaimTemplates = append(sts.Spec.VolumeClaimTemplates, corev1.PersistentVolumeClaim{
			ObjectMeta: metav1.ObjectMeta{
				Namespace: tc.Namespace,
				Name:      "tidb-" + slowlogVolumeName,
			},
		})
	}

	err := tk.cli.Create(context.Background(), sts)
	require.Nil(tk.T, err)
}

func (tk *toolKit) getPod(tc *v1alpha1.TidbCluster, idx int) *corev1.Pod {
	pod := new(corev1.Pod)
	key := types.NamespacedName{Namespace: tc.Namespace, Name: tc.Name + "-tidb-" + strconv.Itoa(idx)}
	err := tk.cli.Get(context.Background(), key, pod)
	if err != nil {
		if apierrors.IsNotFound(err) {
			return nil
		}
		require.Nil(tk.T, err)
	}
	return pod
}

func (tk *toolKit) replacePod(tc *v1alpha1.TidbCluster, idx int, ready bool, withVolume bool) {
	pod := &corev1.Pod{
		ObjectMeta: metav1.ObjectMeta{
			Namespace: tc.Namespace,
			Name:      tc.Name + "-tidb-" + strconv.Itoa(idx),
		},
	}

	if ready {
		pod.Status.Conditions = append(pod.Status.Conditions, corev1.PodCondition{
			Type:   corev1.PodReady,
			Status: corev1.ConditionTrue,
		})
	}

	if withVolume {
		pod.Spec.Volumes = append(pod.Spec.Volumes, corev1.Volume{
			Name: "tidb-" + slowlogVolumeName,
		})
	}

	err := tk.cli.Delete(context.Background(), pod)
	err = client.IgnoreNotFound(err)
	require.Nil(tk.T, err)
	err = tk.cli.Create(context.Background(), pod)
	require.Nil(tk.T, err)
}