package base

import (
	"context"
	"fmt"
	"sync"
	"time"

	"github.com/prometheus/client_golang/prometheus"
	"github.com/samber/lo"
	"github.com/tidbcloud/dataflow-service/app/repo/entity"
	"github.com/tidbcloud/dataflow-service/pkg/log"
	"go.uber.org/zap"
)

type Target[T entity.IdAndVersioning] struct {
	Name                string
	StatusCollector     *prometheus.Desc
	ListUnfinished      func(ctx context.Context) ([]T, error)
	ListByIDs           func(ctx context.Context, ids []entity.ID) ([]T, error)
	UpdateStatusMetrics func(collector *prometheus.Desc, obj T)
	DeleteStatusMetrics func(obj T)
}

const deleteStatusMetricsDelay = 3 * time.Minute

type metricsCollector[T entity.IdAndVersioning] struct {
	target     *Target[T]
	mu         sync.Mutex
	unfinished map[entity.ID]T
}

func NewMetricsCollector[T entity.IdAndVersioning](target *Target[T]) prometheus.Collector {
	return &metricsCollector[T]{
		target:     target,
		unfinished: make(map[entity.ID]T),
	}
}

func (c *metricsCollector[T]) Describe(ch chan<- *prometheus.Desc) {
	ch <- c.target.StatusCollector
}

//nolint:lll
func (c *metricsCollector[T]) Collect(ch chan<- prometheus.Metric) {
	logger := log.GetLogger(context.Background()).Named(fmt.Sprintf("%sMetricsCollector", c.target.Name))
	start := time.Now()
	defer func() {
		logger.Info("metrics collect time", zap.Int64("time in millisecond", time.Since(start).Milliseconds()))
	}()

	ctx := context.Background()

	unfinished, err := c.target.ListUnfinished(ctx)
	if err != nil {
		logger.Error("retrieve unfinished objects failed", zap.Error(err))
		return
	}

	c.mu.Lock()
	defer c.mu.Unlock()

	lastUnfinishedSet := lo.Keys(c.unfinished)
	newUnfinishedSet := lo.Map(unfinished, func(t T, _ int) entity.ID {
		return t.GetID()
	})
	finishedSet, _ := lo.Difference(lastUnfinishedSet, newUnfinishedSet)
	logger.Info("count statistic", zap.Int("unfinished count", len(unfinished)), zap.Int("finished count", len(finishedSet)))

	// update metrics for unfinished objects
	for _, obj := range unfinished {
		c.unfinished[obj.GetID()] = obj
		c.target.UpdateStatusMetrics(c.target.StatusCollector, obj)
	}
	// update and delete metrics for finished objects
	// This keeps terminal status visible for one scrape window before cleanup.
	if len(finishedSet) > 0 {
		// fetch the finished objects to delete their metrics
		finishedObjs, err := c.target.ListByIDs(ctx, finishedSet)
		if err != nil {
			logger.Error("retrieve finished objects failed", zap.Error(err))
			return
		}
		for _, obj := range finishedObjs {
			c.target.UpdateStatusMetrics(c.target.StatusCollector, obj)
			if c.target.DeleteStatusMetrics != nil {
				c.deleteStatusMetricsAsync(obj)
			}
			delete(c.unfinished, obj.GetID())
		}
	}
}

func (c *metricsCollector[T]) deleteStatusMetricsAsync(obj T) {
	go func(deleteStatusMetrics func(T), targetObj T) {
		time.Sleep(deleteStatusMetricsDelay)
		deleteStatusMetrics(targetObj)
	}(c.target.DeleteStatusMetrics, obj)
}