package dm_controller

import (
	"context"
	"testing"
	"time"

	"github.com/prometheus/client_golang/prometheus"
	"github.com/prometheus/client_model/go"
	"github.com/stretchr/testify/require"
	"github.com/tidbcloud/pkgv2/bizerr"
	"github.com/tidbcloud/serverless-service/internal/domain/constant"
	"github.com/tidbcloud/serverless-service/internal/domain/entity"
	"github.com/tidbcloud/serverless-service/internal/domain/repo/mock_repo"
	"go.uber.org/mock/gomock"
)

func TestMigrationCollector_Collect_SkipNilAndExpiredFailedTime(t *testing.T) {
	ctrl := gomock.NewController(t)
	defer ctrl.Finish()

	repo := mock_repo.NewMockDMTaskRepo(ctrl)
	collector := NewMigrationCollector(repo)
	now := time.Now()

	repo.EXPECT().ListTasks(gomock.Any(), gomock.Any()).DoAndReturn(
		func(_ context.Context, filter *entity.DMTaskFilter) ([]*entity.DMTask, bizerr.BizError) {
			if len(filter.Stages) > 0 && filter.Stages[0] == constant.DMTaskStageFailed {
				require.NotNil(t, filter.FailedTimeGTE)
				return []*entity.DMTask{
					{
						ID:         "task-recent",
						ClusterID:  "c1",
						TenantID:   "t1",
						FailedTime: ptrTime(now.Add(-2 * time.Minute)),
						ErrMsg:     "err recent",
						Stage:      constant.DMTaskStageFailed,
					},
				}, nil
			}
			return []*entity.DMTask{}, nil
		},
	).Times(2)
	repo.EXPECT().ListVpcEndpoints(gomock.Any(), gomock.Any()).Return([]*entity.DMVpcEndpoint{}, nil)

	ch := make(chan prometheus.Metric, 16)
	collector.Collect(ch)
	close(ch)

	var failedMetricCount int
	for metric := range ch {
		dtoMetric := &io_prometheus_client.Metric{}
		require.NoError(t, metric.Write(dtoMetric))
		labels := metricLabels(dtoMetric)
		if labels["task_id"] != "" {
			failedMetricCount++
			require.Equal(t, "task-recent", labels["task_id"])
		}
	}
	require.Equal(t, 1, failedMetricCount)
}

func ptrTime(v time.Time) *time.Time {
	return &v
}

func metricLabels(m *io_prometheus_client.Metric) map[string]string {
	labels := make(map[string]string, len(m.GetLabel()))
	for _, label := range m.GetLabel() {
		labels[label.GetName()] = label.GetValue()
	}
	return labels
}