package infra

import (
	"context"
	"testing"

	"github.com/stretchr/testify/require"
	"github.com/tidbcloud/aws-shared-provider/pkg/utils"
	"github.com/tidbcloud/aws-shared-provider/pkg/utils/labels"
	cloudv1 "github.com/tidbcloud/infra-api/api/v1"
	serverlessinfra "github.com/tidbcloud/tidb-management-service/api/serverless_infra"
	infraMeta "github.com/tidbcloud/tidb-management-service/api/serverless_infra/meta"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
)

type testRegionalInfra struct {
	cluster *serverlessinfra.Cluster
	err     error
}

func (f testRegionalInfra) GetCluster(context.Context, string, string) (*serverlessinfra.Cluster, error) {
	return f.cluster, f.err
}

func (f testRegionalInfra) GetClusterWithToken(context.Context, string) (*serverlessinfra.Cluster, error) {
	return f.cluster, f.err
}

type testInfraAPI struct {
	branch *cloudv1.Branch
	err    error
}

func (f testInfraAPI) GetBranch(context.Context, string, string) (*cloudv1.Branch, error) {
	return f.branch, f.err
}

func (f testInfraAPI) GetBranchWithToken(context.Context, string) (*cloudv1.Branch, error) {
	return f.branch, f.err
}

func (f testInfraAPI) PatchBranch(context.Context, string, string, []byte) (*cloudv1.Branch, error) {
	return f.branch, f.err
}

func (testInfraAPI) Start(context.Context) error {
	return nil
}

func (testInfraAPI) WaitForCacheSync(context.Context) bool {
	return true
}

func TestGetClusterOrBranchMetaIncludesInfraTiDBReplicas(t *testing.T) {
	cluster := &serverlessinfra.Cluster{
		ObjectMeta: infraMeta.ObjectMeta{
			TagsKubernetes: "test-kube",
		},
		Status: serverlessinfra.ClusterStatus{
			Phase:    serverlessinfra.ClusterPhaseRunning,
			TidbPool: "default",
			TiDB: serverlessinfra.TiDBStatus{
				UpdatedReplicas: 2,
			},
		},
	}

	branch := &cloudv1.Branch{
		ObjectMeta: metav1.ObjectMeta{
			Labels:      labels.NewKubeClusterLabels(nil, "test-kube"),
			Annotations: map[string]string{utils.TiDBBranchMetadataAnnKey: `{"tenant-id":1,"project-id":2}`},
		},
		Status: cloudv1.BranchStatus{
			Phase: cloudv1.BranchPhaseRunning,
			DB: &cloudv1.DBStatus{
				TiDB: &cloudv1.TiDBStatus{UpdatedReplicas: 3},
			},
		},
	}
	branch.SetAwsFreeTierClusterStatus(&cloudv1.AwsFreeTierClusterStatus{
		TidbPool:    "default",
		RuLimitMode: "standard",
	})

	tests := []struct {
		name          string
		regionalInfra RegionalInfra
		infraAPI      InfraAPI
		wantReplicas  int
	}{
		{
			name:          "cluster",
			regionalInfra: testRegionalInfra{cluster: cluster},
			infraAPI:      testInfraAPI{},
			wantReplicas:  2,
		},
		{
			name:          "branch",
			regionalInfra: testRegionalInfra{err: ErrNotFound},
			infraAPI:      testInfraAPI{branch: branch},
			wantReplicas:  3,
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			meta, err := GetClusterOrBranchMeta(
				context.Background(),
				tt.regionalInfra,
				tt.infraAPI,
				"token",
				"test-kube",
			)
			require.NoError(t, err)
			require.Equal(t, tt.wantReplicas, meta.InfraTiDBReplicas)
		})
	}
}