package premiumdm

import (
	"context"
	"strings"

	"github.com/pingcap/log"
	"github.com/prometheus/client_golang/prometheus"
	"github.com/tidbcloud/dataflow-service-ng/app/models/premium_dm"
	"github.com/tidbcloud/dataflow-service-ng/app/repo/entity"
	"go.uber.org/zap"
)

type StatusCollector struct {
	dao *premium_dm.PremiumDMTaskDAO
}

func NewStatusCollector(dao *premium_dm.PremiumDMTaskDAO) prometheus.Collector {
	return &StatusCollector{dao: dao}
}

func (c *StatusCollector) Describe(ch chan<- *prometheus.Desc) {
	ch <- PremiumDMTaskStatusCount
}

func (c *StatusCollector) Collect(ch chan<- prometheus.Metric) {
	if err := c.collect(context.Background(), ch); err != nil {
		log.L().Error("failed to collect premium dm status metrics", zap.Error(err))
	}
}

func (c *StatusCollector) collect(ctx context.Context, ch chan<- prometheus.Metric) error {
	tasks, err := c.dao.ListTasks(ctx, &entity.PremiumDMTaskFilter{Stages: entity.NonDeletedPremiumDMTaskStages})
	if err != nil {
		return err
	}
	return c.collectFromTasks(tasks, ch)
}

func (c *StatusCollector) collectFromTasks(tasks []*entity.PremiumDMTask, ch chan<- prometheus.Metric) error {
	const statusKeyParts = 3 // tenant/cluster/stage
	statusCount := make(map[string]float64)
	for _, task := range tasks {
		statusKey := strings.Join([]string{task.TenantID, task.TiDBID, string(task.Stage)}, "/")
		statusCount[statusKey]++
	}

	for key, count := range statusCount {
		parts := strings.SplitN(key, "/", statusKeyParts)
		if len(parts) != statusKeyParts {
			continue
		}
		ch <- prometheus.MustNewConstMetric(
			PremiumDMTaskStatusCount,
			prometheus.GaugeValue,
			count,
			parts[0],
			parts[1],
			parts[2],
		)
	}
	return nil
}