package premiumdm

import (
	"context"
	"time"

	"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"
)

const recentFailDuration = 15 * time.Minute

type Collector struct {
	dao *premium_dm.PremiumDMTaskDAO
}

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

func (c *Collector) Describe(ch chan<- *prometheus.Desc) {
	ch <- PremiumDMTaskStageDurationSeconds
	ch <- PremiumDMRecentFail
}

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

func (c *Collector) 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, time.Now(), ch)
}

func (c *Collector) collectFromTasks(tasks []*entity.PremiumDMTask, now time.Time, ch chan<- prometheus.Metric) error {
	for _, task := range tasks {
		stageSince := task.CreatedAt
		if task.UpdatedAt != nil && task.UpdatedAt.After(stageSince) {
			stageSince = *task.UpdatedAt
		}
		ch <- prometheus.MustNewConstMetric(
			PremiumDMTaskStageDurationSeconds,
			prometheus.GaugeValue,
			now.Sub(stageSince).Seconds(),
			task.TenantID,
			task.TiDBID,
			task.ID,
			string(task.Stage),
		)

		if task.Stage != entity.PremiumDMTaskStageFailed {
			continue
		}
		failedAt := stageSince
		if task.FailedTime != nil {
			failedAt = *task.FailedTime
		}
		if now.Sub(failedAt) <= recentFailDuration {
			ch <- prometheus.MustNewConstMetric(
				PremiumDMRecentFail,
				prometheus.GaugeValue,
				now.Sub(failedAt).Seconds(),
				task.TenantID,
				task.TiDBID,
				task.ID,
				string(task.Stage),
			)
		}
	}
	return nil
}