package cluster

import (
	"context"
	"encoding/json"
	"net/http"
	"net/http/httptest"
	"net/url"
	"strconv"
	"strings"
	"testing"
	"time"

	"github.com/go-logr/logr"
	"github.com/gorilla/mux"
	"github.com/stretchr/testify/assert"
	serverlessv1alpha1 "github.com/tidbcloud/aws-shared-provider/apis/serverless/v1alpha1"
	"github.com/tidbcloud/aws-shared-provider/pkg/tidbmanager"
	"github.com/tidbcloud/aws-shared-provider/pkg/utils"
	"github.com/tidbcloud/aws-shared-provider/pkg/utils/clientfake"
	ep "github.com/tidbcloud/aws-shared-provider/pkg/utils/endpoint"
	"github.com/tidbcloud/aws-shared-provider/pkg/utils/pd"
	"github.com/tidbcloud/aws-shared-provider/pkg/utils/query"
	"github.com/tidbcloud/aws-shared-provider/pkg/utils/tidb"
	infra "github.com/tidbcloud/tidb-management-service/api/serverless_infra"
	infraclient "github.com/tidbcloud/tidb-management-service/api/serverless_infra/client"
	"github.com/tidbcloud/tidb-management-service/api/serverless_infra/meta"
	corev1 "k8s.io/api/core/v1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/runtime"
	"k8s.io/client-go/tools/record"
	ctrl "sigs.k8s.io/controller-runtime"
	"sigs.k8s.io/controller-runtime/pkg/client"
	"sigs.k8s.io/controller-runtime/pkg/client/fake"
	"sigs.k8s.io/controller-runtime/pkg/log"

	"github.com/tidbcloud/aws-shared-provider/pkg/scheme"
	"github.com/tidbcloud/aws-shared-provider/pkg/types"
	cloudv1 "github.com/tidbcloud/infra-api/api/v1"
)

const (
	recorderBufSize  = 4096
	defaultNamespace = "cluster-test"
	defaultTidbPool  = "tidb-pool"
)

func TestReconcile(t *testing.T) {
	t.Setenv("NAMESPACE", "tidb-admin")

	deletionTimestamp := meta.Now()
	objs := []runtime.Object{
		&corev1.Namespace{
			ObjectMeta: metav1.ObjectMeta{
				Name: "free-tier-test-network-ns"},
		},
		&corev1.ConfigMap{
			ObjectMeta: metav1.ObjectMeta{
				Name:      types.ServerlessPoolConfig,
				Namespace: defaultTidbPool,
			},
			Data: map[string]string{"tidb-max-replicas": "2"},
		},
		&serverlessv1alpha1.Keyspace{
			ObjectMeta: metav1.ObjectMeta{
				Name:      "kt-test",
				Namespace: "tidb-admin",
			},
			Spec: serverlessv1alpha1.KeyspaceSpec{KeyspaceName: "keyspace-test"},
			Status: serverlessv1alpha1.KeyspaceStatus{
				Phase:      serverlessv1alpha1.AvailablePhase,
				InPool:     true,
				KeyspaceId: 42,
			},
		},
	}

	cases := []struct {
		name string

		objs         []runtime.Object
		infraObjs    []infra.Object
		tidbManagers map[string]tidbmanager.TidbManager

		requeueAfter time.Duration
	}{
		{
			name: "create_cluster",
			objs: []runtime.Object{
				&corev1.Service{
					ObjectMeta: metav1.ObjectMeta{
						Name:      "tidb-gateway",
						Namespace: "tidb-admin",
						Annotations: map[string]string{
							types.AnnoGatewayDnsName: "gateway01.us-east-1.dev.shared.aws.tidbcloud.com",
						},
					},
				},
			},
			infraObjs: []infra.Object{
				&infra.Cluster{
					ObjectMeta: meta.ObjectMeta{
						Name:        "create_cluster",
						Namespace:   defaultNamespace,
						TagsTenant:  "1",
						TagsProject: "1",
					},
					Spec: infra.ClusterSpec{TiDB: infra.TiDBSpec{Replicas: 2}},
				},
			},
		},
		{
			name: "reconcile_cluster",
			objs: []runtime.Object{
				&corev1.Service{
					ObjectMeta: metav1.ObjectMeta{
						Name:      "tidb-gateway",
						Namespace: "tidb-admin",
						Annotations: map[string]string{
							types.AnnoGatewayDnsName: "gateway.tidbcloud.com",
						},
					},
				},
			},
			infraObjs: []infra.Object{
				&infra.Cluster{
					ObjectMeta: meta.ObjectMeta{
						Name:        "reconcile_cluster",
						Namespace:   defaultNamespace,
						TagsTenant:  "1",
						TagsProject: "1",
					},
					Spec: infra.ClusterSpec{TiDB: infra.TiDBSpec{
						Replicas:       2,
						ResourceGroups: map[string]*infra.ResourceGroup{"rg1": {Replicas: 2}},
					}},
					Status: infra.ClusterStatus{
						Phase:          infra.ClusterPhaseRunning,
						TidbPool:       defaultTidbPool,
						TiDB:           infra.TiDBStatus{UpdatedReplicas: 2},
						KeyspaceId:     42,
						ResourceGroups: map[string]infra.ResourceGroup{"rg1": {Pool: defaultTidbPool, Replicas: 2}},
					},
				},
			},
		},
		{
			name: "reconcile_cluster_ensure_resourceGroup",
			objs: []runtime.Object{
				&corev1.Service{
					ObjectMeta: metav1.ObjectMeta{
						Name:      "tidb-gateway",
						Namespace: "tidb-admin",
						Annotations: map[string]string{
							types.AnnoGatewayDnsName: "gateway.tidbcloud.com",
						},
					},
				},
			},
			infraObjs: []infra.Object{
				&infra.Cluster{
					ObjectMeta: meta.ObjectMeta{
						Name:        "reconcile_cluster_ensure_resourceGroup",
						Namespace:   defaultNamespace,
						TagsTenant:  "1",
						TagsProject: "1",
					},
					Spec: infra.ClusterSpec{TiDB: infra.TiDBSpec{
						Replicas: 1,
						ResourceGroups: map[string]*infra.ResourceGroup{
							"rg1": {Replicas: 1},
							"rg3": {Pool: "other-pool", Replicas: 2},
						},
					}},
					Status: infra.ClusterStatus{
						Phase:      infra.ClusterPhaseRunning,
						TidbPool:   defaultTidbPool,
						TiDB:       infra.TiDBStatus{UpdatedReplicas: 1},
						KeyspaceId: 42,
						ResourceGroups: map[string]infra.ResourceGroup{
							"rg1": {Pool: defaultTidbPool, Replicas: 1},
							"rg2": {Pool: "other-pool", Replicas: 1},
						},
					},
				},
			},
			tidbManagers: map[string]tidbmanager.TidbManager{
				defaultTidbPool: &clientfake.FakeTidbManager{},
				"other-pool":    &clientfake.FakeTidbManager{},
			},
		},
		{
			name: "reconcile_cluster_tidb_replicas_changing",
			objs: []runtime.Object{
				&corev1.Service{
					ObjectMeta: metav1.ObjectMeta{
						Name:      "tidb-gateway",
						Namespace: "tidb-admin",
						Annotations: map[string]string{
							types.AnnoGatewayDnsName: "gateway.tidbcloud.com",
						},
					},
				},
			},
			infraObjs: []infra.Object{
				&infra.Cluster{
					ObjectMeta: meta.ObjectMeta{
						Name:        "reconcile_cluster_tidb_replicas_changing",
						Namespace:   defaultNamespace,
						TagsTenant:  "1",
						TagsProject: "1",
					},
					Spec: infra.ClusterSpec{TiDB: infra.TiDBSpec{
						Replicas: 2,
						ResourceGroups: map[string]*infra.ResourceGroup{
							"rg1": {Replicas: 1},
							"rg2": {Pool: "other-pool", Replicas: 2},
						},
					}},
					Status: infra.ClusterStatus{
						Phase:         infra.ClusterPhaseRunning,
						TidbPool:      defaultTidbPool,
						TiDB:          infra.TiDBStatus{UpdatedReplicas: 1},
						KeyspaceToken: "keyspace-token",
						KeyspaceId:    42,
						ResourceGroups: map[string]infra.ResourceGroup{
							"rg1": {Pool: defaultTidbPool, Replicas: 2},
							"rg2": {Pool: "other-pool", Replicas: 1},
						},
					},
				},
			},
			tidbManagers: map[string]tidbmanager.TidbManager{
				defaultTidbPool: &clientfake.FakeTidbManager{
					GetAssignedTidbByTokenR: map[string]*ep.Endpoints{
						"keyspace-token":     {Endpoints: []string{"tidb-server-pod1,ip1"}},
						"keyspace-token/rg1": {Endpoints: []string{"tidb-server-pod2,ip2", "tidb-server-pod3,ip3"}},
					},
				},
				"other-pool": &clientfake.FakeTidbManager{
					GetAssignedTidbByTokenR: map[string]*ep.Endpoints{
						"keyspace-token/rg2": {Endpoints: []string{"tidb-server-pod4,ip4"}},
					},
				},
			},
		},
		{
			name: "reconcile_cluster_AnnoKeyTidbPoolChanging",
			objs: []runtime.Object{
				&corev1.Service{
					ObjectMeta: metav1.ObjectMeta{
						Name:      "tidb-gateway",
						Namespace: "tidb-admin",
						Annotations: map[string]string{
							types.AnnoGatewayDnsName: "gateway.tidbcloud.com",
						},
					},
				},
			},
			infraObjs: []infra.Object{
				&infra.Cluster{
					ObjectMeta: meta.ObjectMeta{
						Name:      "reconcile_cluster_AnnoKeyTidbPoolChanging",
						Namespace: defaultNamespace,
						Annotations: map[string]string{
							types.AnnoKeyTidbPoolChanging: "2024-08-10T00:00:00+00:00",
						},
						TagsTenant:  "1",
						TagsProject: "1",
					},
					Spec: infra.ClusterSpec{TiDB: infra.TiDBSpec{
						Replicas:       2,
						ResourceGroups: map[string]*infra.ResourceGroup{"rg1": {Replicas: 2}},
					}},
					Status: infra.ClusterStatus{
						Phase:          infra.ClusterPhaseRunning,
						TidbPool:       defaultTidbPool,
						TiDB:           infra.TiDBStatus{UpdatedReplicas: 2},
						KeyspaceToken:  "keyspace-token",
						KeyspaceId:     42,
						ResourceGroups: map[string]infra.ResourceGroup{"rg1": {Pool: defaultTidbPool, Replicas: 2}},
					},
				},
			},
			tidbManagers: map[string]tidbmanager.TidbManager{
				defaultTidbPool: &clientfake.FakeTidbManager{
					GetAssignedTidbByTokenR: map[string]*ep.Endpoints{
						"keyspace-token":     {Endpoints: []string{"tidb-server-pod1,ip1"}},
						"keyspace-token/rg1": {Endpoints: []string{"tidb-server-pod2,ip2"}},
					},
				},
				"other-pool": &clientfake.FakeTidbManager{
					GetAssignedTidbByTokenR: map[string]*ep.Endpoints{
						"keyspace-token":     {Endpoints: []string{"tidb-server-pod3,ip3", "tidb-server-pod3,ip3"}},
						"keyspace-token/rg1": {Endpoints: []string{"tidb-server-pod4,ip4", "tidb-server-pod5,ip5"}},
					},
				},
			},
			requeueAfter: retryTidbPoolChanging,
		},
		{
			name: "reconcile_cluster_empty_pool_AnnoKeyTidbPoolChanging",
			objs: []runtime.Object{
				&corev1.Service{
					ObjectMeta: metav1.ObjectMeta{
						Name:      "tidb-gateway",
						Namespace: "tidb-admin",
						Annotations: map[string]string{
							types.AnnoGatewayDnsName: "gateway.tidbcloud.com",
						},
					},
				},
			},
			infraObjs: []infra.Object{
				&infra.Cluster{
					ObjectMeta: meta.ObjectMeta{
						Name:      "reconcile_cluster_empty_pool_AnnoKeyTidbPoolChanging",
						Namespace: defaultNamespace,
						Annotations: map[string]string{
							types.AnnoKeyTidbPoolChanging: time.Now().Format(time.RFC3339),
						},
						TagsTenant:  "1",
						TagsProject: "1",
					},
					Spec: infra.ClusterSpec{TiDB: infra.TiDBSpec{
						Replicas: 2,
						ResourceGroups: map[string]*infra.ResourceGroup{
							"rg1": {Replicas: 1},
							"rg2": {Pool: "other-pool", Replicas: 3},
						},
					}},
					Status: infra.ClusterStatus{
						Phase:         infra.ClusterPhaseRunning,
						TidbPool:      defaultTidbPool,
						TiDB:          infra.TiDBStatus{UpdatedReplicas: 2},
						KeyspaceToken: "keyspace-token",
						KeyspaceId:    42,
						ResourceGroups: map[string]infra.ResourceGroup{
							"rg1": {Pool: defaultTidbPool, Replicas: 1},
							"rg2": {Pool: "other-pool", Replicas: 3},
						},
					},
				},
			},
		},
		{
			name: "delete_cluster",
			infraObjs: []infra.Object{
				&infra.Cluster{
					ObjectMeta: meta.ObjectMeta{
						Name:              "delete_cluster",
						Namespace:         defaultNamespace,
						TagsTenant:        "1",
						TagsProject:       "1",
						DeletionTimestamp: &deletionTimestamp,
					},
					Spec: infra.ClusterSpec{TiDB: infra.TiDBSpec{Replicas: 2}},
					Status: infra.ClusterStatus{
						Phase: infra.ClusterPhaseRunning,
						TiDB:  infra.TiDBStatus{UpdatedReplicas: 2},
					},
				},
			},
		},
	}

	httpServer := newFakeHTTP()
	defer httpServer.Close()

	for _, c := range cases {
		t.Log("test ", c.name)

		r := newFakeReconciler(append(objs, c.objs...), c.infraObjs, httpServer)
		r.TidbManagers = c.tidbManagers
		etcdDB := make(map[string]string)
		for pool, tManager := range c.tidbManagers {
			for keyspaceToken, eps := range tManager.(*clientfake.FakeTidbManager).GetAssignedTidbByTokenR {
				etcdDB[r.EtcdTokenPrefixMap[pool]+keyspaceToken+"/"+utils.KeyReplicas] = strconv.Itoa(len(eps.Endpoints))
				etcdDB[r.EtcdTokenPrefixMap[pool]+keyspaceToken+"/"+utils.KeyEndpoints] = strings.Join(eps.Endpoints, ";")
			}
		}
		r.EtcdKV = clientfake.NewEtcdKV(clientfake.BuildArg{DB: etcdDB})

		req := ctrl.Request{
			NamespacedName: client.ObjectKey{
				Namespace: defaultNamespace,
				Name:      c.name,
			},
		}

		res, err := r.Reconcile(context.Background(), req)
		assert.NoError(t, err)
		assert.Equal(t, c.requeueAfter, res.RequeueAfter)
	}
}

type recordingInfraClient struct {
	infraclient.Client
	consistentGets []bool
}

func (c *recordingInfraClient) Get(ctx context.Context, obj infra.Object, opts ...infraclient.GetOptionFunc) error {
	var option infraclient.GetOption
	for _, opt := range opts {
		opt(&option)
	}
	c.consistentGets = append(c.consistentGets, option.Consistent)
	return c.Client.Get(ctx, obj, opts...)
}

func TestNewReconcileContextConsistentGets(t *testing.T) {
	const kubeClusterName = "test-kubernetes"

	cases := []struct {
		name                   string
		enableInfraMysqlClient bool
		expectedConsistentGets []bool
	}{
		{
			name:                   "mysql client",
			enableInfraMysqlClient: true,
			expectedConsistentGets: []bool{true},
		},
		{
			name:                   "cached client",
			enableInfraMysqlClient: false,
			expectedConsistentGets: []bool{false, true},
		},
	}

	for _, c := range cases {
		t.Run(c.name, func(t *testing.T) {
			httpServer := newFakeHTTP()
			defer httpServer.Close()

			cluster := &infra.Cluster{
				ObjectMeta: meta.ObjectMeta{
					Name:           c.name,
					Namespace:      defaultNamespace,
					TagsKubernetes: kubeClusterName,
				},
				Spec: infra.ClusterSpec{
					TiDB: infra.TiDBSpec{Replicas: 1},
				},
				Status: infra.ClusterStatus{
					TidbPool: defaultTidbPool,
				},
			}
			r := newFakeReconciler([]runtime.Object{
				&corev1.ConfigMap{
					ObjectMeta: metav1.ObjectMeta{
						Name:      types.ServerlessPoolConfig,
						Namespace: defaultTidbPool,
					},
					Data: map[string]string{types.TidbMaxReplicasKey: "2"},
				},
			}, []infra.Object{cluster}, httpServer)
			recordingClient := &recordingInfraClient{Client: r.InfraClusterCli}
			r.InfraClusterCli = recordingClient
			r.EnableInfraMysqlClient = c.enableInfraMysqlClient
			r.KubeClusterName = kubeClusterName

			req := ctrl.Request{
				NamespacedName: client.ObjectKey{
					Namespace: cluster.Namespace,
					Name:      cluster.Name,
				},
			}
			reconcileContext, err := r.newReconcileContext(context.Background(), req)

			assert.NoError(t, err)
			assert.NotNil(t, reconcileContext)
			assert.Equal(t, c.expectedConsistentGets, recordingClient.consistentGets)
		})
	}
}

type fakeTidbController struct {
	tidb.Controller
}

func (tc *fakeTidbController) Stop(ctx context.Context, host, keyspaceName string, graceful bool, waitSeconds int, skipAutoIDOwner, needMgrFree bool) error {
	return nil
}

func newFakeReconciler(objs []runtime.Object, infraObjs []infra.Object, fakeServer *httptest.Server) *Reconciler {
	fakeClient := fake.NewClientBuilder().
		WithScheme(scheme.Scheme).
		WithRuntimeObjects(objs...).
		WithStatusSubresource(&serverlessv1alpha1.Keyspace{}).
		WithIndex(&serverlessv1alpha1.Keyspace{}, "status.phase", func(obj client.Object) []string {
			// filter out BindPhase, it's too more
			if obj.(*serverlessv1alpha1.Keyspace).Status.Phase != serverlessv1alpha1.BindPhase {
				return []string{string(obj.(*serverlessv1alpha1.Keyspace).Status.Phase)}
			}
			return nil
		}).
		WithIndex(&serverlessv1alpha1.Keyspace{}, "status.inPool", func(obj client.Object) []string {
			if obj.(*serverlessv1alpha1.Keyspace).Status.InPool {
				return []string{"true"}
			}
			return nil
		}).
		WithIndex(&serverlessv1alpha1.Keyspace{}, "spec.keyspaceName", func(obj client.Object) []string {
			if keyspaceName := obj.(*serverlessv1alpha1.Keyspace).Spec.KeyspaceName; keyspaceName != "" {
				return []string{keyspaceName}
			}
			return nil
		}).
		WithIndex(&serverlessv1alpha1.Keyspace{}, "status.bindCluster", func(obj client.Object) []string {
			if bindCluster := obj.(*serverlessv1alpha1.Keyspace).Status.BindCluster; bindCluster != "" {
				return []string{bindCluster}
			}
			return nil
		}).
		Build()
	fakeInfraClient := fake.NewClientBuilder().
		WithScheme(scheme.Scheme).
		WithIndex(&cloudv1.Import{}, "spec.clusterName", func(obj client.Object) []string {
			importObj := obj.(*cloudv1.Import)
			return []string{importObj.Spec.ClusterName}
		}).
		Build()

	fakeRecorder := record.NewFakeRecorder(recorderBufSize)
	fakeServerUrl, _ := url.Parse(fakeServer.URL)
	return &Reconciler{
		InfraClusterCli: clientfake.NewInfraClient(infraObjs...).WithPatch(func(set map[string]any, obj infra.Object) error {
			for k, v := range set {
				if anno := strings.TrimPrefix(k, "annotations.'"); anno != "" {
					if obj.(*infra.Cluster).Annotations == nil {
						obj.(*infra.Cluster).Annotations = map[string]string{}
					}
					obj.(*infra.Cluster).Annotations[anno[:len(anno)-1]] = v.(string)
				}
			}
			return nil
		}).Build(),
		Client:      fakeClient,
		Recorder:    fakeRecorder,
		InfraClient: fakeInfraClient,
		TidbManagers: map[string]tidbmanager.TidbManager{
			defaultTidbPool: &clientfake.FakeTidbManager{},
			"other-pool":    &clientfake.FakeTidbManager{},
		},
		EtcdKV:    clientfake.NewEtcdKV(clientfake.BuildArg{}),
		EtcdLease: clientfake.NewEtcdLease(),
		BaseProxy: "fakeproxy.com",
		EtcdTokenPrefixMap: map[string]string{
			defaultTidbPool: "/default-pool/",
			"other-pool":    "/other-pool/",
		},
		TidbPools:      []string{defaultTidbPool, "other-pool"},
		Logger:         logr.New(log.NullLogSink{}),
		PdClient:       pd.NewClient(fakeServerUrl.Host, nil),
		TidbController: &fakeTidbController{},
	}
}

func newFakeReconcileContext(
	cluster *infra.Cluster,
	reconciler *Reconciler,
	httpServer *httptest.Server,
) *reconcileContext {
	if reconciler == nil {
		reconciler = newFakeReconciler(nil, []infra.Object{cluster}, httpServer)
	}

	return &reconcileContext{
		Reconciler:         reconciler,
		ctx:                context.TODO(),
		cluster:            cluster,
		keyspaceToken:      cluster.Status.KeyspaceToken,
		keyspaceID:         cluster.Status.KeyspaceId,
		specTidbPool:       reconciler.getTidbPool(cluster),
		specResourceGroups: cluster.Spec.TiDB.ResourceGroups,
		Logger:             logr.New(log.NullLogSink{}),
	}
}

func newFakeHTTP() *httptest.Server {
	router := mux.NewRouter()
	router.HandleFunc("/pd/api/v2/keyspaces/{keyspace}", func(w http.ResponseWriter, r *http.Request) {
		if mux.Vars(r)["keyspace"] == "" {
			w.WriteHeader(http.StatusBadRequest)
			return
		}
		w.WriteHeader(http.StatusOK)
	}).Methods("GET")
	router.HandleFunc("/pd/api/v2/keyspaces/{keyspace}/config", func(w http.ResponseWriter, r *http.Request) {
		if mux.Vars(r)["keyspace"] == "" {
			w.WriteHeader(http.StatusBadRequest)
			return
		}
		w.WriteHeader(http.StatusOK)
	}).Methods("PATCH")
	router.HandleFunc("/pd/api/v2/keyspaces/{keyspace}/state", func(w http.ResponseWriter, r *http.Request) {
		if mux.Vars(r)["keyspace"] == "" {
			w.WriteHeader(http.StatusBadRequest)
			return
		}
		w.WriteHeader(http.StatusOK)
	}).Methods("PUT")
	router.HandleFunc("/resource-manager/api/v1/config/group", func(w http.ResponseWriter, r *http.Request) {
		w.WriteHeader(http.StatusOK)
	}).Methods("POST", "PUT")
	router.HandleFunc("/resource-manager/api/v1/config/group/{group}", func(w http.ResponseWriter, r *http.Request) {
		if mux.Vars(r)["group"] == "" {
			w.WriteHeader(http.StatusBadRequest)
			return
		}
		w.WriteHeader(http.StatusOK)
	}).Methods("DELETE")
	router.HandleFunc("/resource-manager/api/v1/config/controller/ru-version/{keyspace}", func(w http.ResponseWriter, r *http.Request) {
		w.WriteHeader(http.StatusOK)
	}).Methods("POST")
	router.HandleFunc("/api/v1/query", func(w http.ResponseWriter, r *http.Request) {
		var response query.PromResponse
		_ = json.NewEncoder(w).Encode(response)
		w.WriteHeader(http.StatusOK)
	}).Methods("GET")

	return httptest.NewServer(router)
}