package branch

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"
	"github.com/stretchr/testify/require"
	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"
	"github.com/tidbcloud/aws-shared-provider/pkg/utils/tikv"
	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"
	infra "github.com/tidbcloud/tidb-management-service/api/serverless_infra"
	"github.com/tidbcloud/tidb-management-service/api/serverless_infra/meta"
)

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

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

	deletionTimestamp := metav1.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    []runtime.Object
		tidbManagers map[string]tidbmanager.TidbManager

		requeueAfter time.Duration
	}{
		{
			name: "create_branch",
			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: []runtime.Object{
				&cloudv1.Branch{
					ObjectMeta: metav1.ObjectMeta{
						Name:      "create_cluster",
						Namespace: defaultNamespace,
						Labels: map[string]string{
							"tags.tidbcloud.com/tenant":  "1",
							"tags.tidbcloud.com/project": "1",
						},
					},
					Spec: cloudv1.BranchSpec{
						Cluster: cloudv1.ClusterTopology{
							TiDB: cloudv1.TiDB{
								Replicas:     1,
								InstanceType: "tidb.freetier",
							},
						},
					},
				},
			},
		},
		{
			name: "reconcile_branch",
			objs: []runtime.Object{
				&corev1.Service{
					ObjectMeta: metav1.ObjectMeta{
						Name:      "tidb-gateway",
						Namespace: "tidb-admin",
						Annotations: map[string]string{
							types.AnnoGatewayDnsName: "gateway.tidbcloud.com",
						},
					},
				},
			},
			infraObjs: []runtime.Object{
				&cloudv1.Branch{
					ObjectMeta: metav1.ObjectMeta{
						Name:      "reconcile_cluster",
						Namespace: defaultNamespace,
						Labels: map[string]string{
							"tags.tidbcloud.com/tenant":  "1",
							"tags.tidbcloud.com/project": "1",
						},
					},
					Spec: cloudv1.BranchSpec{
						Cluster: cloudv1.ClusterTopology{
							TiDB: cloudv1.TiDB{
								Replicas:     1,
								InstanceType: "tidb.freetier",
							},
						},
					},
					Status: cloudv1.BranchStatus{
						Phase: cloudv1.BranchPhaseRunning,
						DB:    &cloudv1.DBStatus{TiDB: &cloudv1.TiDBStatus{UpdatedReplicas: 1}},
						FromProvider: runtime.RawExtension{
							Object: &cloudv1.AwsFreeTierClusterStatus{KeyspaceId: 42, TidbPool: defaultTidbPool},
						},
					},
				},
			},
		},
		{
			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: []runtime.Object{
				&cloudv1.Branch{
					ObjectMeta: metav1.ObjectMeta{
						Name:      "reconcile_cluster_tidb_replicas_changing",
						Namespace: defaultNamespace,
						Labels: map[string]string{
							"tags.tidbcloud.com/tenant":  "1",
							"tags.tidbcloud.com/project": "1",
						},
					},
					Spec: cloudv1.BranchSpec{
						Cluster: cloudv1.ClusterTopology{
							TiDB: cloudv1.TiDB{
								Replicas:     1,
								InstanceType: "tidb.freetier",
							},
						},
					},
					Status: cloudv1.BranchStatus{
						Phase: cloudv1.BranchPhaseRunning,
						DB:    &cloudv1.DBStatus{TiDB: &cloudv1.TiDBStatus{UpdatedReplicas: 1}},
						FromProvider: runtime.RawExtension{
							Object: &cloudv1.AwsFreeTierClusterStatus{
								KeyspaceToken: "keyspace-token",
								KeyspaceId:    42,
								TidbPool:      defaultTidbPool,
							},
						},
					},
				},
			},
			tidbManagers: map[string]tidbmanager.TidbManager{
				defaultTidbPool: &clientfake.FakeTidbManager{
					GetAssignedTidbByTokenR: map[string]*ep.Endpoints{
						"keyspace-token": {Endpoints: []string{"tidb-server-pod1,ip1"}},
					},
				},
			},
		},
		{
			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: []runtime.Object{
				&cloudv1.Branch{
					ObjectMeta: metav1.ObjectMeta{
						Name:      "reconcile_cluster_AnnoKeyTidbPoolChanging",
						Namespace: defaultNamespace,
						Annotations: map[string]string{
							types.AnnoKeyTidbPoolChanging: "2024-08-10T00:00:00+00:00",
						},
						Labels: map[string]string{
							"tags.tidbcloud.com/tenant":  "1",
							"tags.tidbcloud.com/project": "1",
						},
					},
					Spec: cloudv1.BranchSpec{
						Cluster: cloudv1.ClusterTopology{
							TiDB: cloudv1.TiDB{
								Replicas:     1,
								InstanceType: "tidb.freetier",
							},
						},
					},
					Status: cloudv1.BranchStatus{
						Phase: cloudv1.BranchPhaseRunning,
						DB:    &cloudv1.DBStatus{TiDB: &cloudv1.TiDBStatus{UpdatedReplicas: 1}},
						FromProvider: runtime.RawExtension{
							Object: &cloudv1.AwsFreeTierClusterStatus{
								KeyspaceToken: "keyspace-token",
								KeyspaceId:    42, TidbPool: defaultTidbPool,
							},
						},
					},
				},
			},
			tidbManagers: map[string]tidbmanager.TidbManager{
				defaultTidbPool: &clientfake.FakeTidbManager{
					GetAssignedTidbByTokenR: map[string]*ep.Endpoints{
						"keyspace-token": {Endpoints: []string{"tidb-server-pod1,ip1"}},
					},
				},
				"other-pool": &clientfake.FakeTidbManager{
					GetAssignedTidbByTokenR: map[string]*ep.Endpoints{
						"keyspace-token": {Endpoints: []string{"tidb-server-pod2,ip2", "tidb-server-pod3,ip3"}},
					},
				},
			},
			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: []runtime.Object{
				&cloudv1.Branch{
					ObjectMeta: metav1.ObjectMeta{
						Name:      "reconcile_cluster_empty_pool_AnnoKeyTidbPoolChanging",
						Namespace: defaultNamespace,
						Annotations: map[string]string{
							types.AnnoKeyTidbPoolChanging: "2024-08-10T00:00:00+00:00",
						},
						Labels: map[string]string{
							"tags.tidbcloud.com/tenant":  "1",
							"tags.tidbcloud.com/project": "1",
						},
					},
					Spec: cloudv1.BranchSpec{
						Cluster: cloudv1.ClusterTopology{
							TiDB: cloudv1.TiDB{
								Replicas:     1,
								InstanceType: "tidb.freetier",
							},
						},
					},
					Status: cloudv1.BranchStatus{
						Phase: cloudv1.BranchPhaseRunning,
						DB:    &cloudv1.DBStatus{TiDB: &cloudv1.TiDBStatus{UpdatedReplicas: 1}},
						FromProvider: runtime.RawExtension{
							Object: &cloudv1.AwsFreeTierClusterStatus{
								KeyspaceToken: "keyspace-token",
								KeyspaceId:    42, TidbPool: defaultTidbPool,
							},
						},
					},
				},
			},
		},
		{
			name: "delete_branch",
			infraObjs: []runtime.Object{
				&cloudv1.Branch{
					ObjectMeta: metav1.ObjectMeta{
						Name:      "delete_cluster",
						Namespace: defaultNamespace,
						Labels: map[string]string{
							"tags.tidbcloud.com/tenant":  "1",
							"tags.tidbcloud.com/project": "1",
						},
						Finalizers:        []string{branchFinalizer},
						DeletionTimestamp: &deletionTimestamp,
					},
					Spec: cloudv1.BranchSpec{
						Cluster: cloudv1.ClusterTopology{
							TiDB: cloudv1.TiDB{
								Replicas:     1,
								InstanceType: "tidb.freetier",
							},
						},
					},
					Status: cloudv1.BranchStatus{
						Phase: cloudv1.BranchPhaseRunning,
						DB:    &cloudv1.DBStatus{TiDB: &cloudv1.TiDBStatus{UpdatedReplicas: 1}},
					},
				},
			},
		},
	}

	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 fakeTidbController struct {
	tidb.Controller
}

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

// fakeTikvApi is a stand-in for the tikv restore API: every restore task
// succeeds immediately, and create calls are counted per mode.
type fakeTikvApi struct {
	createRestoreCalls        int
	createOnlineSnapshotCalls int
	restoreTimepoint          time.Time
	getProgressStatus         string // empty defaults to "Succeed"
}

func (a *fakeTikvApi) GetRestoreProgress(
	_ context.Context, _ uint64, _ uint64, _ string,
) (*tikv.RestoreProgress, error) {
	status := a.getProgressStatus
	if status == "" {
		status = "Succeed"
	}
	return &tikv.RestoreProgress{Status: status}, nil
}

func (a *fakeTikvApi) CreateRestore(
	_ context.Context, _ uint64, _ uint64, _ string, _ string, timepoint time.Time,
) error {
	a.createRestoreCalls++
	a.restoreTimepoint = timepoint
	return nil
}

func (a *fakeTikvApi) CreateOnlineSnapshotRestore(
	_ context.Context, _ uint64, _ uint64, _ string, _ string,
) error {
	a.createOnlineSnapshotCalls++
	return nil
}

func (a *fakeTikvApi) CancelLoadDataTasks(_ context.Context, _ uint64, _ uint32, _ int64) error {
	return nil
}

// newRestoreBranchFixture mirrors the create_branch fixture with cluster
// parent lineage; parentTimestamp controls the ParentTimestampAnnKey
// annotation (empty string means absent).
func newRestoreBranchFixture(parentTimestamp string) *cloudv1.Branch {
	annotations := map[string]string{
		ParentTypeAnnKey:     "cluster",
		ParentBranchIdAnnKey: "parent_cluster",
	}
	if parentTimestamp != "" {
		annotations[ParentTimestampAnnKey] = parentTimestamp
	}
	return &cloudv1.Branch{
		ObjectMeta: metav1.ObjectMeta{
			Name:      "restore_mode_branch",
			Namespace: defaultNamespace,
			Labels: map[string]string{
				"tags.tidbcloud.com/tenant":  "1",
				"tags.tidbcloud.com/project": "1",
			},
			Annotations: annotations,
		},
		Spec: cloudv1.BranchSpec{
			Cluster: cloudv1.ClusterTopology{
				TiDB: cloudv1.TiDB{
					Replicas:     1,
					InstanceType: "tidb.freetier",
				},
			},
		},
	}
}

func TestCreateRestoreKeyspaceOnlineSnapshot(t *testing.T) {
	branch := newRestoreBranchFixture("") // no parent timestamp on purpose
	server := newFakeHTTP()
	reconciler := newFakeReconciler(nil, []runtime.Object{branch}, server)
	reconciler.InfraClusterCli = clientfake.NewInfraClient(&infra.Cluster{
		ObjectMeta: meta.ObjectMeta{Namespace: defaultNamespace, Name: "parent_cluster"},
		Status:     infra.ClusterStatus{Phase: infra.ClusterPhaseRunning, KeyspaceToken: "keyspace-test"},
	}).Build()
	tikvApi := &fakeTikvApi{}
	reconciler.TikvApi = tikvApi
	reconciler.BranchRestoreMode = tikv.RestoreModeOnlineSnapshot
	rx := newFakeReconcileContext(branch, reconciler, server)

	require.NoError(t, rx.createRestoreKeyspace())
	require.Equal(t, 1, tikvApi.createOnlineSnapshotCalls)
	require.Equal(t, 0, tikvApi.createRestoreCalls)
	require.NotEmpty(t, branch.Annotations[types.AnnoKeyBranchRestoreId])
}

// Without an explicit parent timestamp (the branch creator omitted it) pitr
// mode falls back to the CR creation time as the restore point.
func TestCreateRestoreKeyspacePitrFallsBackToCreationTimestamp(t *testing.T) {
	createdAt := time.Now().UTC().Add(-time.Hour).Truncate(time.Second)
	branch := newRestoreBranchFixture("")
	branch.CreationTimestamp = metav1.Time{Time: createdAt}
	server := newFakeHTTP()
	reconciler := newFakeReconciler(nil, []runtime.Object{branch}, server)
	reconciler.InfraClusterCli = clientfake.NewInfraClient(&infra.Cluster{
		ObjectMeta: meta.ObjectMeta{Namespace: defaultNamespace, Name: "parent_cluster"},
		Status:     infra.ClusterStatus{Phase: infra.ClusterPhaseRunning, KeyspaceToken: "keyspace-test"},
	}).Build()
	tikvApi := &fakeTikvApi{}
	reconciler.TikvApi = tikvApi
	reconciler.BranchRestoreMode = tikv.RestoreModePitr
	rx := newFakeReconcileContext(branch, reconciler, server)

	require.NoError(t, rx.createRestoreKeyspace())
	require.Equal(t, 1, tikvApi.createRestoreCalls)
	require.Equal(t, 0, tikvApi.createOnlineSnapshotCalls)
	require.Equal(t, createdAt, tikvApi.restoreTimepoint)
}

func TestCreateRestoreKeyspacePitrUsesTimestamp(t *testing.T) {
	branch := newRestoreBranchFixture("2024-08-10T00:00:00Z")
	server := newFakeHTTP()
	reconciler := newFakeReconciler(nil, []runtime.Object{branch}, server)
	reconciler.InfraClusterCli = clientfake.NewInfraClient(&infra.Cluster{
		ObjectMeta: meta.ObjectMeta{Namespace: defaultNamespace, Name: "parent_cluster"},
		Status:     infra.ClusterStatus{Phase: infra.ClusterPhaseRunning, KeyspaceToken: "keyspace-test"},
	}).Build()
	tikvApi := &fakeTikvApi{}
	reconciler.TikvApi = tikvApi
	reconciler.BranchRestoreMode = tikv.RestoreModePitr
	rx := newFakeReconcileContext(branch, reconciler, server)

	require.NoError(t, rx.createRestoreKeyspace())
	require.Equal(t, 1, tikvApi.createRestoreCalls)
	require.Equal(t, 0, tikvApi.createOnlineSnapshotCalls)
	wantTs, err := time.Parse(time.RFC3339, "2024-08-10T00:00:00Z")
	require.NoError(t, err)
	require.Equal(t, wantTs, tikvApi.restoreTimepoint)
}

// An explicit parent timestamp asks for a restore to that historical point,
// which an online snapshot cannot express; even in online_snapshot mode the
// reconciler must honor it via PiTR.
func TestCreateRestoreKeyspaceOnlineSnapshotHonorsParentTimestamp(t *testing.T) {
	branch := newRestoreBranchFixture("2024-08-10T00:00:00Z")
	server := newFakeHTTP()
	reconciler := newFakeReconciler(nil, []runtime.Object{branch}, server)
	reconciler.InfraClusterCli = clientfake.NewInfraClient(&infra.Cluster{
		ObjectMeta: meta.ObjectMeta{Namespace: defaultNamespace, Name: "parent_cluster"},
		Status:     infra.ClusterStatus{Phase: infra.ClusterPhaseRunning, KeyspaceToken: "keyspace-test"},
	}).Build()
	tikvApi := &fakeTikvApi{}
	reconciler.TikvApi = tikvApi
	reconciler.BranchRestoreMode = tikv.RestoreModeOnlineSnapshot
	rx := newFakeReconcileContext(branch, reconciler, server)

	require.NoError(t, rx.createRestoreKeyspace())
	require.Equal(t, 1, tikvApi.createRestoreCalls)
	require.Equal(t, 0, tikvApi.createOnlineSnapshotCalls)
	wantTs, err := time.Parse(time.RFC3339, "2024-08-10T00:00:00Z")
	require.NoError(t, err)
	require.Equal(t, wantTs, tikvApi.restoreTimepoint)
}

func TestRestorePollInterval(t *testing.T) {
	cases := []struct {
		name          string
		sinceCreation time.Duration
		want          time.Duration
	}{
		{name: "just created", sinceCreation: 0, want: restoreFastPollInterval},
		{name: "inside fast window", sinceCreation: 59 * time.Second, want: restoreFastPollInterval},
		{name: "window boundary", sinceCreation: 60 * time.Second, want: restoreSlowPollInterval},
		{name: "long restore", sinceCreation: 5 * time.Minute, want: restoreSlowPollInterval},
	}
	for _, c := range cases {
		t.Run(c.name, func(t *testing.T) {
			assert.Equal(t, c.want, restorePollInterval(c.sinceCreation))
		})
	}
}

// newRestoringBranchFixture returns a Creating-phase branch mid keyspace
// restore: restore task already created, tidb cluster already ensured,
// progress not yet terminal.
func newRestoringBranchFixture(createdAt time.Time) *cloudv1.Branch {
	return &cloudv1.Branch{
		ObjectMeta: metav1.ObjectMeta{
			Name:      "restoring_branch",
			Namespace: defaultNamespace,
			Labels: map[string]string{
				"tags.tidbcloud.com/tenant":  "1",
				"tags.tidbcloud.com/project": "1",
			},
			Annotations: map[string]string{
				types.AnnoKeyBranchRestoreId: "123456",
			},
			CreationTimestamp: metav1.Time{Time: createdAt},
		},
		Spec: cloudv1.BranchSpec{
			Cluster: cloudv1.ClusterTopology{
				TiDB: cloudv1.TiDB{
					Replicas:     1,
					InstanceType: "tidb.freetier",
				},
			},
		},
		Status: cloudv1.BranchStatus{
			Phase:  cloudv1.BranchPhaseCreating,
			Reason: "still restoring keyspace",
			DB:     &cloudv1.DBStatus{TiDB: &cloudv1.TiDBStatus{UpdatedReplicas: 1}},
			FromProvider: runtime.RawExtension{
				Object: &cloudv1.AwsFreeTierClusterStatus{
					KeyspaceToken: "keyspace-token",
					KeyspaceId:    42,
					TidbPool:      defaultTidbPool,
				},
			},
		},
	}
}

func TestCreateBranchRestorePollInterval(t *testing.T) {
	cases := []struct {
		name         string
		createdAt    time.Time
		wantInterval time.Duration
	}{
		{
			name:         "fresh branch polls fast",
			createdAt:    time.Now().UTC(),
			wantInterval: restoreFastPollInterval,
		},
		{
			name:         "branch older than window polls slow",
			createdAt:    time.Now().UTC().Add(-2 * time.Minute),
			wantInterval: restoreSlowPollInterval,
		},
	}
	for _, c := range cases {
		t.Run(c.name, func(t *testing.T) {
			branch := newRestoringBranchFixture(c.createdAt)
			server := newFakeHTTP()
			reconciler := newFakeReconciler(nil, []runtime.Object{branch}, server)
			reconciler.TikvApi = &fakeTikvApi{getProgressStatus: "Running"}
			rx := newFakeReconcileContext(branch, reconciler, server)

			err := rx.createBranch()
			require.Error(t, err)
			rq, ok := err.(*utils.RequeueAfterError)
			require.True(t, ok, "expected RequeueAfterError, got %T: %v", err, err)
			assert.Equal(t, c.wantInterval, rq.RequeueAfter)
		})
	}
}

func newFakeReconciler(objs, infraObjs []runtime.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).
		WithRuntimeObjects(infraObjs...).
		WithStatusSubresource(&cloudv1.Cluster{}, &cloudv1.Branch{}).
		Build()

	fakeRecorder := record.NewFakeRecorder(recorderBufSize)
	fakeInfraRecorder := record.NewFakeRecorder(recorderBufSize)
	fakeServerUrl, _ := url.Parse(fakeServer.URL)
	return &Reconciler{
		Client:         fakeClient,
		Recorder:       fakeRecorder,
		InfraClient:    fakeInfraClient,
		InfraAPIReader: fakeInfraClient,
		InfraRecorder:  fakeInfraRecorder,
		EtcdKV:         clientfake.NewEtcdKV(clientfake.BuildArg{}),
		EtcdLease:      clientfake.NewEtcdLease(),
		S3BucketName:   "s3_bucket",
		BaseProxy:      "fakeproxy.com",
		ProviderType:   types.ProviderTypeAWSFreeTier,
		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(
	branch *cloudv1.Branch,
	reconciler *Reconciler,
	httpServer *httptest.Server,
) *reconcileContext {
	if reconciler == nil {
		reconciler = newFakeReconciler(nil, []runtime.Object{branch}, httpServer)
	}
	branchStatus, _ := branch.AwsFreeTierClusterStatusFromBranch()

	return &reconcileContext{
		Reconciler:    reconciler,
		ctx:           context.TODO(),
		branch:        branch,
		keyspaceToken: branchStatus.KeyspaceToken,
		keyspaceID:    branchStatus.KeyspaceId,
		specTidbPool:  reconciler.getTidbPool(branch),
		Logger:        logr.New(log.NullLogSink{}),
	}
}

func newFakeHTTP() *httptest.Server {
	router := mux.NewRouter()
	router.HandleFunc("/pd/api/v1/cluster", func(w http.ResponseWriter, r *http.Request) {
		_, _ = w.Write([]byte(`{"id": 1}`))
	}).Methods("GET")
	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("/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)
}