package premiumdm

import (
	"testing"
	"time"

	"github.com/prometheus/client_golang/prometheus"
	io_prometheus_client "github.com/prometheus/client_model/go"
	"github.com/stretchr/testify/require"
	"github.com/tidbcloud/dataflow-service-ng/app/repo/entity"
)

func TestCollectorEmitsPremiumDMMetrics(t *testing.T) {
	now := time.Now()
	updatedAt := now.Add(-10 * time.Minute)
	failedAt := now.Add(-5 * time.Minute)

	tasks := []*entity.PremiumDMTask{
		{
			ID:        "task-running",
			TenantID:  "1001",
			TiDBID:    "2001",
			Stage:     entity.PremiumDMTaskStageRunning,
			CreatedAt: now.Add(-20 * time.Minute),
			UpdatedAt: &updatedAt,
		},
		{
			ID:         "task-failed",
			TenantID:   "1001",
			TiDBID:     "2001",
			Stage:      entity.PremiumDMTaskStageFailed,
			CreatedAt:  now.Add(-30 * time.Minute),
			UpdatedAt:  &updatedAt,
			FailedTime: &failedAt,
		},
	}

	collector := &Collector{dao: nil}
	ch := make(chan prometheus.Metric, 16)
	require.NoError(t, collector.collectFromTasks(tasks, now, ch))
	close(ch)

	var recentFailFound, runningDurationFound bool
	for metric := range ch {
		m := &io_prometheus_client.Metric{}
		require.NoError(t, metric.Write(m))
		labels := map[string]string{}
		for _, lp := range m.GetLabel() {
			labels[lp.GetName()] = lp.GetValue()
		}
		switch {
		case labels["stage"] == string(entity.PremiumDMTaskStageRunning) && labels["task_id"] == "task-running":
			runningDurationFound = true
			require.GreaterOrEqual(t, m.GetGauge().GetValue(), 600.0)
		case labels["stage"] == string(entity.PremiumDMTaskStageFailed) && labels["task_id"] == "task-failed":
			recentFailFound = true
			require.Greater(t, m.GetGauge().GetValue(), 0.0)
		}
	}

	require.True(t, recentFailFound)
	require.True(t, runningDurationFound)
}

func TestStatusCollectorEmitsPremiumDMStatusCount(t *testing.T) {
	tasks := []*entity.PremiumDMTask{
		{TenantID: "1001", TiDBID: "2001", Stage: entity.PremiumDMTaskStageRunning},
		{TenantID: "1001", TiDBID: "2001", Stage: entity.PremiumDMTaskStageRunning},
		{TenantID: "1001", TiDBID: "2001", Stage: entity.PremiumDMTaskStageFailed},
	}

	collector := &StatusCollector{dao: nil}
	ch := make(chan prometheus.Metric, 16)
	require.NoError(t, collector.collectFromTasks(tasks, ch))
	close(ch)

	statusCounts := map[string]float64{}
	for metric := range ch {
		m := &io_prometheus_client.Metric{}
		require.NoError(t, metric.Write(m))
		labels := map[string]string{}
		for _, lp := range m.GetLabel() {
			labels[lp.GetName()] = lp.GetValue()
		}
		statusCounts[labels["stage"]] = m.GetGauge().GetValue()
	}

	require.Equal(t, 2.0, statusCounts[string(entity.PremiumDMTaskStageRunning)])
	require.Equal(t, 1.0, statusCounts[string(entity.PremiumDMTaskStageFailed)])
}