package dmscaler

import (
	"context"
	"crypto/tls"
	"fmt"
	"sync"
	"time"

	"github.com/go-logr/logr"
	"github.com/pingcap/tidb-operator/pkg/apis/pingcap/v1alpha1"
	"github.com/tidbcloud/aws-shared-provider/pkg/utils/dm"
	dmapi "github.com/tidbcloud/pkgv2/clusterproxy/serverless/dm"
	"github.com/tidbcloud/pkgv2/log"
	"github.com/tidbcloud/pkgv2/worker"
	"go.uber.org/zap"
	"k8s.io/apimachinery/pkg/types"
	ctrl "sigs.k8s.io/controller-runtime"
	"sigs.k8s.io/controller-runtime/pkg/client"
)

const (
	// Constants for better maintainability
	defaultNamespace  = "tidb-serverless"
	scaleLoopInterval = 5 * time.Second
	scaleTaskTimeout  = 2 * time.Minute

	scaleDownCooldown = 1 * time.Minute

	// Replica limits
	minReplicas = 0    // Minimum number of worker replicas
	maxReplicas = 1000 // Maximum number of worker replicas

	dmServiceSuffix = "-dm-master.tidb-serverless:8261"
)

type DMWorkerScaler struct {
	logr.Logger

	tlsConfig *tls.Config
	client    client.Client
	wg        sync.WaitGroup
	namespace string
}

func New(client client.Client, tlsConfig *tls.Config) *DMWorkerScaler {
	scaler := &DMWorkerScaler{
		Logger:    ctrl.Log.WithName("dmworkerscaler"),
		tlsConfig: tlsConfig,
		client:    client,
		namespace: defaultNamespace,
	}

	return scaler
}

func (s *DMWorkerScaler) Start(ctx context.Context) error {
	// List all DMClusters in the namespace
	dmClusters, err := s.listDMClusters(ctx)
	if err != nil {
		return fmt.Errorf("failed to list DMClusters: %w", err)
	}

	// If no clusters found, return early
	if len(dmClusters) == 0 {
		s.Info("No DMClusters found, exiting", "namespace", s.namespace)
		return nil
	}

	s.Info("Found DMClusters, starting scale loops", "count", len(dmClusters), "namespace", s.namespace)

	// Start a scale loop for each cluster
	for _, cluster := range dmClusters {
		s.wg.Add(1)
		api := dm.NewApi(cluster.Name+dmServiceSuffix, s.tlsConfig)
		scaler := &clusterScaler{
			Logger:      s.Logger.WithValues("dmCluster", cluster.Name),
			api:         api,
			client:      s.client,
			namespace:   s.namespace,
			clusterName: cluster.Name,
		}
		// Start the scale loop in a separate goroutine
		go scaler.scaleLoopForCluster(ctx, &s.wg)
	}

	s.wg.Wait()
	return nil
}

func (s *DMWorkerScaler) NeedLeaderElection() bool {
	return true
}

func (s *DMWorkerScaler) listDMClusters(ctx context.Context) ([]v1alpha1.DMCluster, error) {
	var dmClusters v1alpha1.DMClusterList

	if err := s.client.List(ctx, &dmClusters, client.InNamespace(s.namespace)); err != nil {
		return nil, fmt.Errorf("failed to list DMClusters: %w", err)
	}

	return dmClusters.Items, nil
}

type clusterScaler struct {
	logr.Logger

	api           dm.Api
	client        client.Client
	namespace     string
	clusterName   string
	lastScaleTime time.Time
}

func (s *clusterScaler) scaleLoopForCluster(ctx context.Context, wg *sync.WaitGroup) {
	defer wg.Done()

	s.Info("Starting DM Worker Scaler for cluster",
		"namespace", s.namespace,
		"clusterName", s.clusterName)

	logger := log.With(zap.String("worker", "dmClusterScaler"), zap.String("interval", scaleLoopInterval.String()))
	periodWorker := worker.NewPeriodWorker(logger, scaleLoopInterval, func(ctx context.Context) error {
		timeoutCtx, cancel := context.WithTimeout(ctx, scaleTaskTimeout)
		defer cancel()
		startTime := time.Now()
		err := s.performScaleTaskForCluster(timeoutCtx)
		if err != nil {
			s.Error(err, "Failed to perform scale task", "clusterName", s.clusterName)
			return err
		}
		s.Info("Scale task completed", "duration", time.Since(startTime), "clusterName", s.clusterName)
		return nil
	})

	periodWorker.Start()
	<-ctx.Done()
	periodWorker.Stop()
	s.Info("Scale loop stopped due to context cancellation", "clusterName", s.clusterName)
}

func (s *clusterScaler) performScaleTaskForCluster(ctx context.Context) error {
	s.Info("Scale task started", "clusterName", s.clusterName)
	// Step 1: Count sources
	sourceCount, err := s.countSources(ctx)
	if err != nil {
		return fmt.Errorf("failed to count sources: %w", err)
	}

	// Step 2: Update DMCluster worker replicas
	if err := s.updateDMClusterReplicasForCluster(ctx, sourceCount); err != nil {
		return fmt.Errorf("failed to update DMCluster replicas: %w", err)
	}

	return nil
}

func (s *clusterScaler) countSources(ctx context.Context) (int, error) {
	var sources dmapi.GetSourceListResponse
	var err error

	sources, err = s.api.ListSources(ctx)
	if err != nil {
		return 0, fmt.Errorf("failed to list sources: %w", err)
	}

	s.Info("Counted sources", "total", len(sources.Data))
	return len(sources.Data), nil
}

func (s *clusterScaler) updateDMClusterReplicasForCluster(ctx context.Context, enabledSourceCount int) error {
	// Get the DMCluster
	dmCluster := &v1alpha1.DMCluster{}
	if err := s.client.Get(ctx, types.NamespacedName{
		Name:      s.clusterName,
		Namespace: s.namespace,
	}, dmCluster); err != nil {
		return fmt.Errorf("failed to get DMCluster %s/%s: %w", s.namespace, s.clusterName, err)
	}

	currentReplicas := int(dmCluster.Spec.Worker.Replicas)
	targetReplicas := enabledSourceCount

	// Enforce replica limits
	if targetReplicas < minReplicas {
		targetReplicas = minReplicas
	} else if targetReplicas > maxReplicas {
		targetReplicas = maxReplicas
	}

	if targetReplicas != currentReplicas {
		operation := "scale-up"
		if targetReplicas < currentReplicas {
			operation = "scale-down"

			// Apply cooldown for scale down operations
			if !s.lastScaleTime.IsZero() && time.Since(s.lastScaleTime) < scaleDownCooldown {
				s.Info("Scale down operation skipped due to cooldown",
					"clusterName", s.clusterName,
					"current", currentReplicas,
					"target", targetReplicas,
					"timeSinceLastScale", time.Since(s.lastScaleTime),
					"cooldown", scaleDownCooldown)
				return nil
			}
		}

		s.Info("Updating DMCluster worker replicas",
			"clusterName", s.clusterName,
			"current", currentReplicas,
			"enabledSources", enabledSourceCount,
			"target", targetReplicas,
			"operation", operation)

		// Use JSON Patch to precisely update only the Worker.Replicas field
		// This is because the dmCluster object may not have the latest CRD version,
		// and using client.Update() would overwrite newly added fields on the server
		// with empty values. JSON Patch only modifies the specified field,
		// avoiding impact on other fields' latest values.
		patchData := fmt.Sprintf(`[{"op": "replace", "path": "/spec/worker/replicas", "value": %d}]`, targetReplicas)
		patch := client.RawPatch(types.JSONPatchType, []byte(patchData))

		if err := s.client.Patch(ctx, dmCluster, patch); err != nil {
			return fmt.Errorf("failed to update DMCluster replicas: %w", err)
		}

		s.lastScaleTime = time.Now()

		s.Info("Successfully updated DMCluster worker replicas",
			"clusterName", s.clusterName,
			"replicas", targetReplicas,
			"operation", operation)
	} else {
		s.Info("No replica adjustment needed",
			"clusterName", s.clusterName,
			"current", currentReplicas,
			"enabledSources", enabledSourceCount)
	}

	return nil
}