package cron

import (
	"context"
	"errors"
	"fmt"
	"math"
	"sync"
	"time"

	apiv1 "github.com/tidbcloud/cluster-service-ng/api/cluster-service-ng/v1beta"
	"github.com/tidbcloud/cluster-service-ng/domain/entity"
	"github.com/tidbcloud/cluster-service-ng/domain/server"
	"github.com/tidbcloud/cluster-service-ng/pkg/clustergateway"
	"github.com/tidbcloud/cluster-service-ng/pkg/storage/repo"
	"github.com/tidbcloud/cluster-service-ng/pkg/utils"
	"github.com/tidbcloud/tiinf/monitor/tilog"
	"go.opentelemetry.io/otel"
	"go.uber.org/zap"
	"golang.org/x/sync/errgroup"
	"google.golang.org/grpc/codes"
	"google.golang.org/grpc/status"
)

const (
	dmWorkerAutoscalerConcurrency = 5

	dmWorkerScaleDirectionUp   = "up"
	dmWorkerScaleDirectionDown = "down"
)

var dmWorkerAutoscalerTracer = otel.Tracer("DMWorkerAutoscalerJob")

type dmWorkerScaleDownCandidate struct {
	target          uint32
	firstObservedAt time.Time
}

// DMWorkerAutoscalerJob keeps DM worker replicas aligned with enabled DM sources.
type DMWorkerAutoscalerJob struct {
	dmClusterRepo      repo.DMClusterRepo
	sharedPoolRepo     repo.SharedPoolRepo
	lockRepo           repo.LockRepo
	ngService          server.Service
	clusterGateway     clustergateway.Client
	expression         string
	scaleDownStability time.Duration

	running             sync.Mutex
	scaleDownCandidates sync.Map
	now                 func() time.Time
}

func NewDMWorkerAutoscalerJob(
	dmClusterRepo repo.DMClusterRepo,
	sharedPoolRepo repo.SharedPoolRepo,
	lockRepo repo.LockRepo,
	ngService server.Service,
	clusterGateway clustergateway.Client,
	expression string,
	scaleDownStability time.Duration,
) *DMWorkerAutoscalerJob {
	return &DMWorkerAutoscalerJob{
		dmClusterRepo:      dmClusterRepo,
		sharedPoolRepo:     sharedPoolRepo,
		lockRepo:           lockRepo,
		ngService:          ngService,
		clusterGateway:     clusterGateway,
		expression:         expression,
		scaleDownStability: scaleDownStability,
		now:                time.Now,
	}
}

func (j *DMWorkerAutoscalerJob) Name() string {
	return "DMWorkerAutoscalerJob"
}

func (j *DMWorkerAutoscalerJob) CronExpression() string {
	return j.expression
}

func (j *DMWorkerAutoscalerJob) Execute(ctx context.Context) error {
	if !j.running.TryLock() {
		tilog.Warn(ctx, "dm worker autoscaler is still running, skipping this execution")
		return nil
	}
	defer j.running.Unlock()

	dmClusters, err := j.dmClusterRepo.List(ctx, &repo.DMClusterListFilter{})
	if err != nil {
		return fmt.Errorf("list DM clusters: %w", err)
	}

	var (
		group    errgroup.Group
		errs     []error
		errsLock sync.Mutex
	)
	group.SetLimit(dmWorkerAutoscalerConcurrency)

	for _, dmCluster := range dmClusters {
		if dmCluster.GetPhase() != apiv1.DMCluster_AVAILABLE {
			j.scaleDownCandidates.Delete(dmCluster.GetId())
			continue
		}

		dmCluster := dmCluster
		group.Go(func() error {
			if err := j.processDMCluster(ctx, dmCluster); err != nil {
				j.scaleDownCandidates.Delete(dmCluster.GetId())
				errsLock.Lock()
				errs = append(errs, fmt.Errorf("DM cluster %s in shared pool %s: %w",
					dmCluster.GetId(), dmCluster.GetSharedPoolId(), err))
				errsLock.Unlock()
			}
			return nil
		})
	}
	_ = group.Wait()

	return errors.Join(errs...)
}

func (j *DMWorkerAutoscalerJob) processDMCluster(ctx context.Context, dmCluster *entity.DMCluster) error {
	sharedPool, err := j.sharedPoolRepo.Load(ctx, dmCluster.GetSharedPoolId(), false)
	if err != nil {
		return fmt.Errorf("load shared pool: %w", err)
	}

	activeSourceCount, err := j.clusterGateway.ListActiveDMSourceCount(ctx, sharedPool, dmCluster.GetId())
	if err != nil {
		return fmt.Errorf("list active DM sources: %w", err)
	}
	if activeSourceCount < 0 || uint64(activeSourceCount) > math.MaxUint32 {
		return fmt.Errorf("invalid active DM source count: %d", activeSourceCount)
	}

	workerQuota := dmCluster.GetWorkerNodeQuota()
	if workerQuota == nil {
		return errors.New("DM worker node quota is missing")
	}

	target := uint32(activeSourceCount)
	current := workerQuota.GetNodeCount()
	switch {
	case target > current:
		j.scaleDownCandidates.Delete(dmCluster.GetId())
		_, err := j.scaleDMWorkers(ctx, dmCluster, target, activeSourceCount, dmWorkerScaleDirectionUp)
		return err
	case target == current:
		j.scaleDownCandidates.Delete(dmCluster.GetId())
		return nil
	default:
		return j.handleScaleDown(ctx, dmCluster, target, activeSourceCount, current)
	}
}

func (j *DMWorkerAutoscalerJob) handleScaleDown(
	ctx context.Context,
	dmCluster *entity.DMCluster,
	target uint32,
	activeSourceCount int,
	current uint32,
) error {
	if j.scaleDownStability <= 0 {
		_, err := j.scaleDMWorkers(ctx, dmCluster, target, activeSourceCount, dmWorkerScaleDirectionDown)
		return err
	}

	now := j.now()
	value, loaded := j.scaleDownCandidates.Load(dmCluster.GetId())
	if !loaded {
		j.scaleDownCandidates.Store(dmCluster.GetId(), dmWorkerScaleDownCandidate{
			target:          target,
			firstObservedAt: now,
		})
		tilog.Info(ctx, "DM worker scale down candidate observed, waiting for stability",
			zap.String("sharedPoolId", dmCluster.GetSharedPoolId()),
			zap.String("dmClusterId", dmCluster.GetId()),
			zap.Int("activeSourceCount", activeSourceCount),
			zap.Uint32("currentWorkers", current),
			zap.Uint32("targetWorkers", target),
			zap.Duration("stabilityDuration", j.scaleDownStability))
		return nil
	}

	candidate := value.(dmWorkerScaleDownCandidate)
	if candidate.target != target {
		j.scaleDownCandidates.Store(dmCluster.GetId(), dmWorkerScaleDownCandidate{
			target:          target,
			firstObservedAt: now,
		})
		tilog.Info(ctx, "DM worker scale down candidate reset",
			zap.String("sharedPoolId", dmCluster.GetSharedPoolId()),
			zap.String("dmClusterId", dmCluster.GetId()),
			zap.Int("activeSourceCount", activeSourceCount),
			zap.Uint32("currentWorkers", current),
			zap.Uint32("previousTargetWorkers", candidate.target),
			zap.Uint32("targetWorkers", target),
			zap.Duration("stabilityDuration", j.scaleDownStability))
		return nil
	}

	if now.Sub(candidate.firstObservedAt) < j.scaleDownStability {
		return nil
	}

	tilog.Info(ctx, "DM worker scale down stability elapsed, scaling down",
		zap.String("sharedPoolId", dmCluster.GetSharedPoolId()),
		zap.String("dmClusterId", dmCluster.GetId()),
		zap.Int("activeSourceCount", activeSourceCount),
		zap.Uint32("currentWorkers", current),
		zap.Uint32("targetWorkers", target),
		zap.Duration("observedFor", now.Sub(candidate.firstObservedAt)))

	completed, err := j.scaleDMWorkers(ctx, dmCluster, target, activeSourceCount, dmWorkerScaleDirectionDown)
	if err != nil {
		return err
	}
	if completed {
		j.scaleDownCandidates.Delete(dmCluster.GetId())
	}
	return nil
}

func (j *DMWorkerAutoscalerJob) scaleDMWorkers(
	ctx context.Context,
	dmCluster *entity.DMCluster,
	target uint32,
	activeSourceCount int,
	direction string,
) (bool, error) {
	unlock, err := j.lockRepo.Lock(ctx, &entity.Lock{
		Key:         "dm_worker_autoscaler_" + utils.FnvHash(dmCluster.GetId()),
		TimeoutSecs: 0,
	})
	if err != nil {
		if status.Code(err) == codes.FailedPrecondition {
			tilog.Info(ctx, "DM worker autoscale skipped, lock held by another worker",
				zap.String("sharedPoolId", dmCluster.GetSharedPoolId()),
				zap.String("dmClusterId", dmCluster.GetId()))
			return false, nil
		}
		return false, fmt.Errorf("acquire DM worker autoscaler lock: %w", err)
	}
	defer func() { _ = unlock() }()

	latest, err := j.ngService.GetDMCluster(ctx, dmCluster.GetSharedPoolId(), dmCluster.GetId())
	if err != nil {
		return false, fmt.Errorf("refresh DM cluster: %w", err)
	}
	if latest.GetPhase() != apiv1.DMCluster_AVAILABLE {
		tilog.Warn(ctx, "DM worker autoscale skipped, cluster not available",
			zap.String("sharedPoolId", dmCluster.GetSharedPoolId()),
			zap.String("dmClusterId", dmCluster.GetId()),
			zap.String("phase", latest.GetPhase().String()),
			zap.Uint32("targetWorkers", target))
		return false, nil
	}

	latestWorkerQuota := latest.GetWorkerNodeQuota()
	if latestWorkerQuota == nil {
		return false, errors.New("latest DM worker node quota is missing")
	}
	latestCurrent := latestWorkerQuota.GetNodeCount()
	if latestCurrent == target {
		tilog.Debug(ctx, "DM worker autoscale skipped, already at target",
			zap.String("sharedPoolId", dmCluster.GetSharedPoolId()),
			zap.String("dmClusterId", dmCluster.GetId()),
			zap.Int("activeSourceCount", activeSourceCount),
			zap.Uint32("targetWorkers", target))
		return true, nil
	}
	if latestCurrent > target && direction != dmWorkerScaleDirectionDown {
		tilog.Info(ctx, "DM worker autoscale skipped, latest state requires scale down stability",
			zap.String("sharedPoolId", dmCluster.GetSharedPoolId()),
			zap.String("dmClusterId", dmCluster.GetId()),
			zap.Int("activeSourceCount", activeSourceCount),
			zap.Uint32("currentWorkers", latestCurrent),
			zap.Uint32("targetWorkers", target),
			zap.String("requestedDirection", direction))
		return false, nil
	}
	if latestCurrent < target && direction == dmWorkerScaleDirectionDown {
		direction = dmWorkerScaleDirectionUp
	}

	// Cron executions do not always have an incoming request trace; create one for DAG id uniqueness.
	scaleCtx, span := dmWorkerAutoscalerTracer.Start(ctx, "ScaleDMCluster")
	defer span.End()

	_, err = j.ngService.ScaleDMCluster(
		scaleCtx,
		dmCluster.GetSharedPoolId(),
		dmCluster.GetId(),
		nil,
		&entity.NodeQuota{NodeCount: target},
		false,
	)
	if err != nil {
		return false, fmt.Errorf("scale DM workers to %d: %w", target, err)
	}

	tilog.Info(ctx, "submitted DM worker autoscale request",
		zap.String("sharedPoolId", dmCluster.GetSharedPoolId()),
		zap.String("dmClusterId", dmCluster.GetId()),
		zap.Int("activeSourceCount", activeSourceCount),
		zap.Uint32("currentWorkers", latestCurrent),
		zap.Uint32("targetWorkers", target),
		zap.String("direction", direction))
	return true, nil
}