package bgtask

import (
	"container/list"
	"context"
	"encoding/json"
	"fmt"
	"net/http"
	"sync"
	"time"

	"github.com/go-logr/logr"
	"github.com/gorilla/mux"
	"github.com/pingcap/errors"
	"github.com/pingcap/kvproto/pkg/keyspacepb"
	"github.com/tidbcloud/aws-shared-provider/pkg/tidbworker/admin"
	"github.com/tidbcloud/aws-shared-provider/pkg/tidbworker/job"
	"github.com/tidbcloud/aws-shared-provider/pkg/tidbworker/schema/bgtask"
	"github.com/tidbcloud/aws-shared-provider/pkg/tidbworker/types"
	types2 "github.com/tidbcloud/aws-shared-provider/pkg/types"
	pd "github.com/tikv/pd/client"
	"gorm.io/gorm"
	"gorm.io/gorm/clause"
	batchv1 "k8s.io/api/batch/v1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/labels"
	"k8s.io/client-go/informers"
	"k8s.io/client-go/kubernetes"
	v1 "k8s.io/client-go/listers/batch/v1"
	"k8s.io/client-go/tools/cache"
	"sigs.k8s.io/controller-runtime/pkg/client"
	ctrl "sigs.k8s.io/controller-runtime/pkg/manager"
)

var _ types.Handler = &handler{}

const (
	pollInterval              = 15 * time.Second
	recycleDoneJobsInterval   = 30 * time.Second
	recycleOrphanJobsInterval = 5 * time.Minute
	updateMetricsInterval     = 15 * time.Second
	autoScaleInterval         = 10 * time.Second
)

const (
	// taskCountHistoryLen is the length of task count history.
	taskCountHistoryLen = 60
	// expectedTaskFinishMinute is the expected time for a task to finish in minutes.
	expectedTaskFinishMinute = 10
)

type taskCount struct {
	count int
	time  time.Time
}

type WorkerContext struct {
	KeyspaceName     string
	TaskCountHistory list.List
}

type handler struct {
	clusterCli       client.Client
	clientset        *kubernetes.Clientset
	namespace        string
	logger           logr.Logger
	stop             chan struct{}
	wg               sync.WaitGroup
	db               *gorm.DB
	pdCli            pd.Client
	jobBuilder       *job.Builder
	adminConfig      *admin.ConfigReconciler
	workerCertSecret string
	batchJobLimit    int
	ddlJobLimit      int
	batchWorkerCount int
	ddlWorkerCount   int

	importIntoJobLimit    int
	importIntoWorkerCount int

	enableWorkerAutoScale          bool
	disableWorkerAutoScaleTaskType map[string]bool

	maxSharedWorkerCount            int
	maxSharedWorkerCountPerKeyspace int
	componentNode                   string
	jobLister                       v1.JobLister
	workersMu                       sync.RWMutex
	workers                         map[string]*WorkerContext
}

func (h *handler) NeedLeaderElection() bool {
	return true
}

func (h *handler) Start(ctx context.Context) error {
	h.logger.Info("starting handler")

	err := h.db.WithContext(ctx).AutoMigrate(&bgtask.Task{}, &bgtask.TaskDone{})
	if err != nil {
		return err
	}

	initMetrics()

	h.wg.Add(4) //nolint:mnd // number of loops
	go h.pollLoop(ctx)
	go h.recycleDoneJobsLoop(ctx)
	go h.recycleOrphanJobsLoop(ctx)
	go h.updateMetricsLoop(ctx)
	if h.enableWorkerAutoScale {
		h.wg.Add(1)
		go h.autoScaleLoop(ctx)
	}
	h.wg.Wait()

	close(h.stop)

	return nil
}

func (h *handler) pollLoop(ctx context.Context) {
	defer h.wg.Done()

	ticker := time.NewTicker(pollInterval)
	defer ticker.Stop()
	for {
		select {
		case <-ctx.Done():
			return
		case <-ticker.C:
			err := h.poll(ctx)
			if err != nil {
				h.logger.Error(err, "failed to poll tasks")
			}
		}
	}
}

func (h *handler) poll(ctx context.Context) error { //nolint:funlen // function is long but clear
	db := h.db.WithContext(ctx)

	batchJobQuota, err := job.CheckJobQuota(
		h.clusterCli, h.namespace, types2.JobTypeLabelKey, string(job.TypeBatch), h.batchJobLimit)
	if err != nil {
		h.logger.Error(err, "failed to count running jobs")
		return err
	}
	ddlJobQuota, err := job.CheckJobQuota(
		h.clusterCli, h.namespace, types2.JobTypeLabelKey, string(job.TypeDDL), h.ddlJobLimit)
	if err != nil {
		h.logger.Error(err, "failed to count running jobs")
		return err
	}
	importIntoQuota, err := job.CheckJobQuota(
		h.clusterCli, h.namespace, types2.JobTypeLabelKey, string(job.TypeImportInto), h.importIntoJobLimit)
	if err != nil {
		h.logger.Error(err, "failed to count running jobs")
		return err
	}
	// Calculate how many jobs we can create.
	remainingJobsQuota := map[job.Type]int{
		job.TypeBatch:      batchJobQuota,
		job.TypeDDL:        ddlJobQuota,
		job.TypeImportInto: importIntoQuota,
	}
	if remainingJobsQuota[job.TypeBatch]+remainingJobsQuota[job.TypeDDL]+remainingJobsQuota[job.TypeImportInto] <= 0 {
		h.logger.Info("job limit reached, skip poll loop",
			"batch-job-limit", h.batchJobLimit,
			"ddl-job-limit", h.ddlJobLimit,
			"import-into-job-limit", h.importIntoJobLimit)
		return nil
	}

	// Acquire pending tasks.
	var pendingTasks []struct {
		KeyspaceName string
		TaskType     string
		TidbPool     string
		GlobalTaskID int64
		ExecID       string
		RegisteredAt time.Time
	}
	err = db.
		Model(&bgtask.Task{}).
		Select("keyspace_name, task_type, tidb_pool, global_task_id, exec_id, MIN(registered_at) as registered_at").
		Where("state = ?", bgtask.StatePending).
		Where("exec_id != ?", "").
		Group("keyspace_name").
		Group("task_type").
		Group("tidb_pool").
		Group("global_task_id").
		Group("exec_id").
		Order("registered_at ASC"). // Handle tasks in FIFO order.
		Find(&pendingTasks).Error
	if err != nil {
		h.logger.Error(err, "failed to query pending tasks")
		return err
	}

	disabled := make(map[string]struct{})
	for _, task := range pendingTasks {
		if _, ok := disabled[task.KeyspaceName]; ok {
			continue
		}
		if remainingJobsQuota[job.Type(task.TaskType)] <= 0 {
			continue
		}

		meta, err := h.pdCli.LoadKeyspace(ctx, task.KeyspaceName)
		if err != nil {
			h.logger.Error(err, "failed to load keyspace", "keyspace-name", task.KeyspaceName)
			continue
		}

		// If the keyspace is not enabled, we need to mark the task as done.
		if meta.State != keyspacepb.KeyspaceState_ENABLED {
			h.logger.Info("skip activating job, keyspace is not enabled", "keyspace-name", task.KeyspaceName)
			disabled[task.KeyspaceName] = struct{}{}
			continue
		}

		if task.TaskType != string(job.TypeDDL) &&
			task.TaskType != string(job.TypeBatch) &&
			task.TaskType != string(job.TypeImportInto) {
			h.logger.Error(errors.New("unknown task type"), "unknown task type", "task-type", task.TaskType)
			continue
		}

		if h.adminConfig != nil && h.adminConfig.IsPaused(task.TaskType) {
			h.logger.Info("skip activating job, task type is paused", "task-type", task.TaskType)
			continue
		}

		if h.enableWorkerAutoScale && !h.disableWorkerAutoScaleTaskType[task.TaskType] {
			// update state to running directly, defer to `autoScale` to create job.
			err := h.markTaskRunning(ctx, task.KeyspaceName, task.ExecID)
			if err != nil {
				h.logger.Error(err, "failed to mark task as running", "keyspace-name", task.KeyspaceName)
			}
			continue
		}

		go func(
			keyspaceName string,
			tidbPool string,
			taskType job.Type,
			globalTaskID int64,
			execID string,
			registeredAt time.Time,
		) {
			h.logger.Info("activating job", "keyspace-name", keyspaceName)
			err = h.activateJob(ctx, keyspaceName, meta.Id, tidbPool, taskType, globalTaskID, execID)
			if err != nil {
				h.logger.Error(err, "failed to activate job", "keyspace-name", keyspaceName)
				return
			}
			runnerCounter[taskType].Inc()
			taskQueueDuration[taskType].Observe(time.Since(registeredAt).Seconds())
		}(task.KeyspaceName, task.TidbPool, job.Type(task.TaskType), task.GlobalTaskID, task.ExecID, task.RegisteredAt)

		remainingJobsQuota[job.Type(task.TaskType)]--
		if remainingJobsQuota[job.Type(task.TaskType)] <= 0 {
			h.logger.Info("job limit reached, skip activation", "task-type", task.TaskType)
			// If both batch and ddl jobs are skipped, we can break the loop.
			if remainingJobsQuota[job.TypeBatch] <= 0 &&
				remainingJobsQuota[job.TypeDDL] <= 0 &&
				remainingJobsQuota[job.TypeImportInto] <= 0 {
				break
			}
			continue
		}
	}
	for k := range disabled {
		err := db.Model(&bgtask.Task{}).
			Where("keyspace_name = ?", k).
			Update("state", bgtask.StateDone).Error
		if err != nil {
			h.logger.Error(err,
				"failed to update disabled keyspaces' task state", "keyspace-name", k)
			return err
		}
	}
	return nil
}

func (h *handler) recycleDoneJobsLoop(ctx context.Context) {
	defer h.wg.Done()

	ticker := time.NewTicker(recycleDoneJobsInterval)
	defer ticker.Stop()
	for {
		select {
		case <-ctx.Done():
			return
		case <-ticker.C:
			err := h.recycleDoneJobs(ctx)
			if err != nil {
				h.logger.Error(err, "failed to recycle done tasks")
			}
		}
	}
}

// recycleDoneJobs terminate runners that have finished their tasks.
// A keyspace's runner is considered only if all tasks belong to the keyspace are done.
func (h *handler) recycleDoneJobs(ctx context.Context) error {
	type DoneGlobalTask struct {
		KeyspaceName string
		GlobalTaskID int64
	}
	var (
		db              = h.db.WithContext(ctx)
		doneGlobalTasks []DoneGlobalTask
		doneSubTasks    []bgtask.Task
	)

	// Find all done global tasks.
	err := db.Model(&bgtask.Task{}).
		Select("keyspace_name, global_task_id").
		Group("keyspace_name, global_task_id").
		Having("COUNT(*) = COUNT(CASE WHEN state = 'done' THEN 1 END)").
		Find(&doneGlobalTasks).Error
	if err != nil {
		h.logger.Error(err, "failed to query done global tasks")
		return err
	}

	// Find all completed subTasks.
	for _, globalTask := range doneGlobalTasks {
		var subTasks []bgtask.Task
		db.Model(&bgtask.Task{}).
			Where("keyspace_name = ? AND global_task_id = ?", globalTask.KeyspaceName, globalTask.GlobalTaskID).
			Find(&subTasks)
		doneSubTasks = append(doneSubTasks, subTasks...)
	}

	// deletedRunners is used to mark the job that has been terminated.
	deletedRunners := map[string]struct{}{"": {}}
	for _, doneTask := range doneSubTasks {
		// Delete target runners.
		// If the job is terminated, ignore this record.
		if !h.enableWorkerAutoScale || h.disableWorkerAutoScaleTaskType[doneTask.TaskType] {
			// if auto scale is enabled, use `autoScale` to cleanup jobs.
			if _, ok := deletedRunners[doneTask.Runner]; !ok {
				j := &batchv1.Job{
					ObjectMeta: metav1.ObjectMeta{
						Name:      doneTask.Runner,
						Namespace: h.namespace,
					},
				}
				h.logger.Info("deleting job", "job", j.Name)
				err := h.clusterCli.Delete(ctx, j, client.PropagationPolicy(metav1.DeletePropagationBackground))
				if client.IgnoreNotFound(err) != nil {
					h.logger.Error(err, "failed to delete job", "job", j.Name)
					continue
				}
			}
		}

		// Delete the task record, then insert the record into done table.
		h.logger.Info("deleting task record", "done-task", doneTask)
		err = db.Transaction(func(tx *gorm.DB) error {
			err = tx.Delete(doneTask).Error
			if err != nil {
				h.logger.Error(err, "failed to delete task record", "done-task", doneTask)
				return err
			}

			err := tx.Clauses(clause.OnConflict{DoNothing: true}).Create(&bgtask.TaskDone{
				KeyspaceName: doneTask.KeyspaceName,
				TaskKey:      doneTask.TaskKey,
				GlobalTaskID: doneTask.GlobalTaskID,
				SubTaskID:    doneTask.SubTaskID,
				TaskType:     doneTask.TaskType,
				TidbPool:     doneTask.TidbPool,
				ExecID:       doneTask.ExecID,
				Runner:       doneTask.Runner,
				State:        doneTask.State,
				RegisteredAt: doneTask.RegisteredAt,
			}).Error
			if err != nil {
				h.logger.Error(err, "failed to create done task record", "done-task", doneTask)
				return err
			}

			return nil
		})
		if err != nil {
			continue
		}
		// Mark the job as terminated.
		deletedRunners[doneTask.Runner] = struct{}{}
		taskCompleteDuration[job.Type(doneTask.TaskType)].Observe(time.Since(doneTask.RegisteredAt).Seconds())
	}
	return nil
}

// recycleOrphanJobsLoop terminate runners belongs to keyspace that are no longer enabled.
func (h *handler) recycleOrphanJobsLoop(ctx context.Context) {
	defer h.wg.Done()

	ticker := time.NewTicker(recycleOrphanJobsInterval)
	defer ticker.Stop()

	for {
		select {
		case <-ctx.Done():
			return
		case <-ticker.C:
			err := h.recycleOrphanJobs(ctx)
			if err != nil {
				h.logger.Error(err, "failed to recycle orphan jobs")
			}
		}
	}
}

func (h *handler) recycleOrphanJobs(ctx context.Context) error {
	var (
		db      = h.db.WithContext(ctx)
		results []struct {
			KeyspaceName string
			Runner       string
		}
	)

	err := db.Model(&bgtask.Task{}).
		Distinct("keyspace_name, runner").
		Where("state = ?", bgtask.StateRunning).
		Find(&results).Error
	if err != nil {
		h.logger.Error(err, "failed to query running tasks")
		return err
	}

	for _, r := range results {
		m, err := h.pdCli.LoadKeyspace(ctx, r.KeyspaceName)
		if err != nil {
			h.logger.Error(err, "failed to load keyspace", "keyspace-name", r.KeyspaceName)
			continue
		}

		if m.State != keyspacepb.KeyspaceState_ENABLED {
			err := db.Model(&bgtask.Task{}).
				Where("keyspace_name = ?", r.KeyspaceName).
				Where("runner = ?", r.Runner).
				Update("state", bgtask.StateDone).Error
			if err != nil {
				h.logger.Error(err,
					"failed to update task state", "keyspace-name", r.KeyspaceName, "runner", r.Runner)
				continue
			}
		}
	}

	return nil
}

func (h *handler) autoScaleLoop(ctx context.Context) {
	defer h.wg.Done()

	ticker := time.NewTicker(autoScaleInterval)
	defer ticker.Stop()

	for {
		select {
		case <-ctx.Done():
			return
		case <-ticker.C:
			err := h.autoScale(ctx)
			if err != nil {
				h.logger.Error(err, "failed to auto scale")
			}
		}
	}
}

func makeSharedExecID(index int) string {
	return fmt.Sprintf("tidb-worker-shared:%d", index)
}

func (h *handler) sharedWorkerCountPerKeyspaceLimit() int {
	if h.maxSharedWorkerCountPerKeyspace > 0 {
		return h.maxSharedWorkerCountPerKeyspace
	}
	return h.maxSharedWorkerCount
}

func (h *handler) limitSharedWorkerCountPerKeyspace(keyspace string, count int) int {
	limit := h.sharedWorkerCountPerKeyspaceLimit()
	if limit <= 0 || count <= limit {
		return count
	}
	h.logger.Info("shared worker count reached per-keyspace limit",
		"keyspace-name", keyspace,
		"need-worker-count", count,
		"max-shared-worker-count-per-keyspace", limit)
	return limit
}

func (h *handler) exceedsSharedWorkerCountPerKeyspaceLimit(count int) bool {
	limit := h.sharedWorkerCountPerKeyspaceLimit()
	return limit > 0 && count > limit
}

func (h *handler) autoScale(ctx context.Context) error {
	db := h.db.WithContext(ctx)
	type Result struct {
		KeyspaceName string
		TidbPool     string
		TaskType     string
		TaskCount    int
	}
	var results []Result
	err := db.
		Model(&bgtask.Task{}).
		Select("keyspace_name, tidb_pool, task_type, COUNT(*) as task_count").
		Where("state != ?", bgtask.StateDone).
		Group("keyspace_name, tidb_pool, task_type").
		Find(&results).Error
	if err != nil {
		h.logger.Error(err, "failed to query task count")
		return err
	}

	type TaskCount struct {
		KeyspaceName string
		KeyspaceID   uint32
		TidbPool     string
		TaskCount    int
	}

	taskCounts := make(map[string]*TaskCount)

	for _, r := range results {
		meta, err := h.pdCli.LoadKeyspace(ctx, r.KeyspaceName)
		if err != nil {
			h.logger.Error(err, "failed to load keyspace", "keyspace-name", r.KeyspaceName)
			continue
		}

		// If the keyspace is not enabled, we need to mark the task as done.
		if meta.State != keyspacepb.KeyspaceState_ENABLED {
			h.logger.Info("skip activating job, keyspace is not enabled", "keyspace-name", r.KeyspaceName)
			continue
		}

		if h.disableWorkerAutoScaleTaskType[r.TaskType] {
			continue
		}

		if _, ok := taskCounts[r.KeyspaceName]; ok {
			taskCounts[r.KeyspaceName].TaskCount += r.TaskCount
		} else {
			taskCounts[r.KeyspaceName] = &TaskCount{
				KeyspaceName: r.KeyspaceName,
				KeyspaceID:   meta.Id,
				TidbPool:     r.TidbPool,
				TaskCount:    r.TaskCount,
			}
		}
	}

	for _, counts := range taskCounts {
		h.UpdateWorkers(ctx, counts.KeyspaceName, counts.KeyspaceID, counts.TidbPool, counts.TaskCount)
	}

	// recycle jobs
	keyspaces := make(map[string]struct{})
	for _, r := range results {
		if h.disableWorkerAutoScaleTaskType[r.TaskType] {
			continue
		}
		keyspaces[r.KeyspaceName] = struct{}{}
	}

	var deleteJobs []*batchv1.Job
	jobs, err := h.getAllJobs()
	if err != nil {
		h.logger.Error(err, "failed to list all workers")
		return err
	}
	for _, job := range jobs {
		if _, ok := keyspaces[job.Labels[types2.KeyspaceNameLabelKey]]; !ok {
			deleteJobs = append(deleteJobs, job)
		}
	}

	h.workersMu.Lock()
	for keyspace := range h.workers {
		if _, ok := keyspaces[keyspace]; !ok {
			delete(h.workers, keyspace)
		}
	}
	h.workersMu.Unlock()

	for _, job := range deleteJobs {
		h.logger.Info("deleting job", "job", job.Name)
		err := h.clusterCli.Delete(ctx, job, client.PropagationPolicy(metav1.DeletePropagationBackground))
		if client.IgnoreNotFound(err) != nil {
			h.logger.Error(err, "failed to delete job", "job", job.Name)
			continue
		}
	}

	return nil
}

//nolint:lll // ignore
func (h *handler) UpdateWorkers(ctx context.Context, keyspace string, keyspaceID uint32, tidbPool string, taskCount int) {
	workers, err := h.getJobs(keyspace)
	if err != nil {
		h.logger.Error(err, "failed to get workers", "keyspace-name", keyspace)
		return
	}
	workerIndics := make(map[int]*batchv1.Job)
	for _, worker := range workers {
		index, err := getWorkerIndexFromJobName(worker.Name)
		if err != nil {
			h.logger.Error(err, "failed to parse worker index", "job-name", worker.Name)
			continue
		}
		workerIndics[index] = worker
	}

	needWorkerCount := 1
	if len(workers) > 0 {
		h.recordTaskCount(keyspace, taskCount)

		// check if more workers are needed
		speed := h.getTaskConsumeSpeed(keyspace)
		h.logger.Info("task consume speed", "keyspace-name", keyspace, "speed", speed, "task-count", taskCount)
		if speed == 0 {
			// Not enough history to scale out yet, but still enforce the configured upper bound.
			needWorkerCount = len(workerIndics)
		} else {
			needWorkerCount = int(float64(taskCount)/speed/expectedTaskFinishMinute) + 1
		}
	}
	needWorkerCount = h.limitSharedWorkerCountPerKeyspace(keyspace, needWorkerCount)

	for index := 0; index < needWorkerCount; index++ {
		if _, ok := workerIndics[index]; ok {
			continue
		}
		execID := makeSharedExecID(index)
		go func(keyspaceName string, keyspaceID uint32, tidbPool string, execID string) {
			h.logger.Info("activating job", "keyspace-name", keyspaceName)
			err = h.activateJobV2(ctx, keyspaceName, keyspaceID, tidbPool, execID)
			if err != nil {
				h.logger.Error(err, "failed to activate job", "keyspace-name", keyspaceName)
				return
			}
			runnerCounter[job.TypeShared].Inc()
		}(keyspace, keyspaceID, tidbPool, execID)
	}

	// try scale in
	if taskCount < len(workerIndics) || h.exceedsSharedWorkerCountPerKeyspaceLimit(len(workerIndics)) {
		// recycle jobs that do not have any tasks
		type Result struct {
			ExecID    string
			TaskCount int
		}
		var results []Result
		err := h.db.WithContext(ctx).
			Model(&bgtask.Task{}).
			Select("exec_id, COUNT(*) as task_count").
			Where("keyspace_name = ? AND state != ?", keyspace, bgtask.StateDone).
			Group("exec_id").
			Find(&results).Error
		if err != nil {
			h.logger.Error(err, "failed to query task count")
			return
		}
		taskCounts := make(map[string]int)
		for _, r := range results {
			taskCounts[r.ExecID] = r.TaskCount
		}
		var maxIndex int
		for index := range workerIndics {
			if index > maxIndex {
				maxIndex = index
			}
		}
		for index := maxIndex; index >= needWorkerCount; index-- {
			if taskCounts[makeSharedExecID(index)] > 0 {
				// If there are still tasks assigned to this worker, stop scaling in further.
				h.logger.Info("cannot scale in further, worker still has tasks", "worker-index", index)
				break
			}

			job := workerIndics[index]
			if job == nil {
				continue
			}
			h.logger.Info("scale in worker, deleting job", "job", job.Name)
			err := h.clusterCli.Delete(ctx, job, client.PropagationPolicy(metav1.DeletePropagationBackground))
			if client.IgnoreNotFound(err) != nil {
				h.logger.Error(err, "failed to delete job", "job", job.Name)
			}
		}
	}
}

func (h *handler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	h.logger.Info("received request", "method", r.Method, "url", r.URL.Path)
	router := mux.NewRouter()
	router.HandleFunc("/scaler/api/v1/bgtask/batch-worker", h.handleBatchWorker)
	router.HandleFunc("/scaler/api/v1/bgtask/ddl-worker", h.handleDDLWorker)
	router.HandleFunc("/scaler/api/v1/bgtask/import-into-worker", h.handleImportIntoWorker)
	router.HandleFunc("/scaler/api/v1/bgtask/batch-worker/{keyspace}", h.handleBatchWorkerV2)
	router.HandleFunc("/scaler/api/v1/bgtask/ddl-worker/{keyspace}", h.handleDDLWorkerV2)
	router.HandleFunc("/scaler/api/v1/bgtask/import-into-worker/{keyspace}", h.handleImportIntoWorkerV2)
	router.HandleFunc("/scaler/api/v1/bgtask/auto-analyze-worker/{keyspace}", h.handleAutoAnalyzeWorkerV2)
	router.ServeHTTP(w, r)
}

func (h *handler) handleBatchWorker(w http.ResponseWriter, r *http.Request) {
	type Response struct {
		Paused      bool `json:"paused"`
		WorkerCount int  `json:"worker-count"`
	}

	var res Response
	if h.adminConfig.IsPaused(string(job.TypeBatch)) {
		res.Paused = true
	} else {
		res.WorkerCount = h.batchWorkerCount
	}

	err := json.NewEncoder(w).Encode(res)
	if err != nil {
		h.logger.Error(err, "failed to encode response")
	}
}

func (h *handler) getAllJobs() ([]*batchv1.Job, error) {
	selector := labels.Set(map[string]string{
		types2.JobTypeLabelKey: string(job.TypeShared),
	}).AsSelector()
	jobs, err := h.jobLister.List(selector)
	if err != nil {
		return nil, err
	}
	return jobs, nil
}

func (h *handler) getJobs(keyspace string) ([]*batchv1.Job, error) {
	selector := labels.Set(map[string]string{
		types2.KeyspaceNameLabelKey: keyspace,
		types2.JobTypeLabelKey:      string(job.TypeShared),
	}).AsSelector()
	jobs, err := h.jobLister.List(selector)
	if err != nil {
		return nil, err
	}
	return jobs, nil
}

// the jobs must have index from 0~N-1
func (h *handler) getExecJobsCount(keyspace string) (int, error) {
	jobs, err := h.getJobs(keyspace)
	if err != nil {
		return 0, err
	}
	jobIndex := make(map[int]struct{})
	for _, job := range jobs {
		index, err := getWorkerIndexFromJobName(job.Name)
		if err != nil {
			h.logger.Error(err, "failed to parse job index", "job-name", job.Name)
			continue
		}
		jobIndex[index] = struct{}{}
	}
	for i := 0; i < len(jobIndex); i++ {
		if _, ok := jobIndex[i]; !ok {
			return i, nil
		}
	}
	return len(jobIndex), nil
}

func (h *handler) recordTaskCount(keyspace string, count int) {
	h.workersMu.Lock()
	defer h.workersMu.Unlock()
	worker := h.workers[keyspace]
	if worker == nil {
		worker = &WorkerContext{
			KeyspaceName: keyspace,
		}
		h.workers[keyspace] = worker
	}
	worker.TaskCountHistory.PushBack(taskCount{count: count, time: time.Now()})
	if worker.TaskCountHistory.Len() > taskCountHistoryLen {
		worker.TaskCountHistory.Remove(worker.TaskCountHistory.Front())
	}
}

// getTaskConsumeSpeed returns the average task consume speed of the keyspace.
// count per minute.
func (h *handler) getTaskConsumeSpeed(keyspace string) float64 {
	h.workersMu.RLock()
	defer h.workersMu.RUnlock()
	worker := h.workers[keyspace]
	if worker == nil || worker.TaskCountHistory.Len() < 2 {
		return 0
	}

	var totalConsume float64
	for e := worker.TaskCountHistory.Front(); e != nil; e = e.Next() {
		if next := e.Next(); next != nil {
			if e.Value.(taskCount).count > next.Value.(taskCount).count {
				totalConsume += float64(e.Value.(taskCount).count - next.Value.(taskCount).count)
			}
		}
	}
	start := worker.TaskCountHistory.Front().Value.(taskCount).time
	end := worker.TaskCountHistory.Back().Value.(taskCount).time
	elapsed := end.Sub(start).Minutes()
	if elapsed <= 0 {
		return 0
	}
	if elapsed >= 1 && totalConsume == 0 {
		totalConsume = 1
	}
	return totalConsume / elapsed
}

func (h *handler) handleBatchWorkerV2(w http.ResponseWriter, r *http.Request) {
	type Response struct {
		Paused          bool `json:"paused"`
		WorkerCount     int  `json:"worker-count"`
		EnableAutoScale bool `json:"enable-auto-scale"`
	}

	var res Response
	if h.adminConfig.IsPaused(string(job.TypeBatch)) {
		res.Paused = true
	} else {
		if h.enableWorkerAutoScale && !h.disableWorkerAutoScaleTaskType[string(job.TypeBatch)] {
			res.EnableAutoScale = true
			count, _ := h.getExecJobsCount(mux.Vars(r)["keyspace"]) // ignore error
			res.WorkerCount = count
		} else {
			res.WorkerCount = h.batchWorkerCount
		}
	}

	err := json.NewEncoder(w).Encode(res)
	if err != nil {
		h.logger.Error(err, "failed to encode response")
	}
}

func (h *handler) handleImportIntoWorker(w http.ResponseWriter, r *http.Request) {
	type Response struct {
		Paused      bool `json:"paused"`
		WorkerCount int  `json:"worker-count"`
	}

	var res Response
	if h.adminConfig.IsPaused(string(job.TypeImportInto)) {
		res.Paused = true
	} else {
		res.WorkerCount = h.importIntoWorkerCount
	}

	err := json.NewEncoder(w).Encode(res)
	if err != nil {
		h.logger.Error(err, "failed to encode response")
	}
}

func (h *handler) handleImportIntoWorkerV2(w http.ResponseWriter, r *http.Request) {
	type Response struct {
		Paused          bool `json:"paused"`
		WorkerCount     int  `json:"worker-count"`
		EnableAutoScale bool `json:"enable-auto-scale"`
	}

	var res Response
	if h.adminConfig.IsPaused(string(job.TypeImportInto)) {
		res.Paused = true
	} else {
		if h.enableWorkerAutoScale && !h.disableWorkerAutoScaleTaskType[string(job.TypeImportInto)] {
			res.EnableAutoScale = true
			count, _ := h.getExecJobsCount(mux.Vars(r)["keyspace"]) // ignore error
			res.WorkerCount = count
		} else {
			res.WorkerCount = h.importIntoWorkerCount
		}
	}

	err := json.NewEncoder(w).Encode(res)
	if err != nil {
		h.logger.Error(err, "failed to encode response")
	}
}

func (h *handler) handleDDLWorker(w http.ResponseWriter, r *http.Request) {
	type Response struct {
		Paused      bool `json:"paused"`
		WorkerCount int  `json:"worker-count"`
	}

	var res Response
	if h.adminConfig.IsPaused(string(job.TypeDDL)) {
		res.Paused = true
	} else {
		res.WorkerCount = h.ddlWorkerCount
	}

	err := json.NewEncoder(w).Encode(res)
	if err != nil {
		h.logger.Error(err, "failed to encode response")
	}
}

func (h *handler) handleDDLWorkerV2(w http.ResponseWriter, r *http.Request) {
	type Response struct {
		Paused          bool `json:"paused"`
		WorkerCount     int  `json:"worker-count"`
		EnableAutoScale bool `json:"enable-auto-scale"`
	}

	var res Response
	if h.adminConfig.IsPaused(string(job.TypeDDL)) {
		res.Paused = true
	} else {
		if h.enableWorkerAutoScale && !h.disableWorkerAutoScaleTaskType[string(job.TypeDDL)] {
			res.EnableAutoScale = true
			count, _ := h.getExecJobsCount(mux.Vars(r)["keyspace"]) // ignore error
			res.WorkerCount = count
		} else {
			res.WorkerCount = h.ddlWorkerCount
		}
	}

	err := json.NewEncoder(w).Encode(res)
	if err != nil {
		h.logger.Error(err, "failed to encode response")
	}
}

func (h *handler) handleAutoAnalyzeWorkerV2(w http.ResponseWriter, r *http.Request) {
	type Response struct {
		Paused          bool `json:"paused"`
		WorkerCount     int  `json:"worker-count"`
		EnableAutoScale bool `json:"enable-auto-scale"`
	}

	var res Response
	if h.adminConfig.IsPaused(string(job.TypeAutoAnalyze)) {
		res.Paused = true
	} else if h.enableWorkerAutoScale && !h.disableWorkerAutoScaleTaskType[string(job.TypeDDL)] {
		res.EnableAutoScale = true
		// worker count does not matter for auto analyze
	}

	err := json.NewEncoder(w).Encode(res)
	if err != nil {
		h.logger.Error(err, "failed to encode response")
	}
}

func (h *handler) updateMetricsLoop(ctx context.Context) {
	defer h.wg.Done()

	ticker := time.NewTicker(updateMetricsInterval)
	defer ticker.Stop()

	for {
		select {
		case <-ctx.Done():
			return
		case <-ticker.C:
			h.updateMetrics(ctx)
		}
	}
}

func (h *handler) updateMetrics(ctx context.Context) {
	var (
		db = h.db.WithContext(ctx)
	)

	type TaskCount struct {
		State    bgtask.State
		TaskType job.Type
		Count    int64
	}

	var taskCounts []TaskCount
	err := db.Model(&bgtask.Task{}).
		Select("state, task_type, count(*) as count").
		Group("state, task_type").
		Find(&taskCounts).Error
	if err != nil {
		h.logger.Error(err, "failed to query task counts")
		return
	}

	countMap := map[job.Type]map[bgtask.State]float64{
		job.TypeBatch: {
			bgtask.StatePending: 0,
			bgtask.StateRunning: 0,
			bgtask.StateDone:    0,
		},
		job.TypeDDL: {
			bgtask.StatePending: 0,
			bgtask.StateRunning: 0,
			bgtask.StateDone:    0,
		},
		job.TypeImportInto: {
			bgtask.StatePending: 0,
			bgtask.StateRunning: 0,
			bgtask.StateDone:    0,
		},
		job.TypeAutoAnalyze: {
			bgtask.StatePending: 0,
			bgtask.StateRunning: 0,
			bgtask.StateDone:    0,
		},
	}

	// Update task count map.
	for _, taskCount := range taskCounts {
		countMap[taskCount.TaskType][taskCount.State] = float64(taskCount.Count)
	}

	for taskType, states := range countMap {
		for state, count := range states {
			currentTaskCounter[taskType][state].Set(count)
		}
	}
}

func NewHandler(
	namespace string,
	mgr ctrl.Manager,
	db *gorm.DB,
	pdCli pd.Client,
	ops ...option) types.Handler {
	h := &handler{
		clusterCli: mgr.GetClient(),
		clientset:  kubernetes.NewForConfigOrDie(mgr.GetConfig()),
		namespace:  namespace,
		logger:     mgr.GetLogger().WithName("bgtask"),
		stop:       make(chan struct{}),
		db:         db,
		pdCli:      pdCli,
		jobBuilder: job.NewBuilder(namespace, mgr.GetClient()),
		workers:    make(map[string]*WorkerContext),

		disableWorkerAutoScaleTaskType: make(map[string]bool),
	}
	for _, op := range ops {
		op(h)
	}
	if h.enableWorkerAutoScale {
		h.logger.Info("auto scale is enabled")

		factory := informers.NewSharedInformerFactoryWithOptions(
			h.clientset, 10*time.Minute, informers.WithNamespace(h.namespace),
		)
		factory.Batch().V1().Jobs().Informer().AddEventHandler(&jobEventHandler{logger: h.logger}) //nolint:errcheck // not need
		h.jobLister = factory.Batch().V1().Jobs().Lister()

		h.logger.Info("create informer factory and lister", "lister", h.jobLister)

		factory.Start(h.stop)
		factory.WaitForCacheSync(h.stop)

		h.logger.Info("informer cache sync done")
	}
	return h
}

type option func(*handler)

func WithAdminConfig(adminConfig *admin.ConfigReconciler) option {
	return func(h *handler) {
		h.adminConfig = adminConfig
	}
}

func WithWorkerCertSecret(workerCertSecret string) option {
	return func(h *handler) {
		h.workerCertSecret = workerCertSecret
	}
}

func WithBatchJobLimit(limit int) option {
	return func(h *handler) {
		h.batchJobLimit = limit
	}
}

func WithDDLJobLimit(limit int) option {
	return func(h *handler) {
		h.ddlJobLimit = limit
	}
}

func WithImportIntoJobLimit(limit int) option {
	return func(h *handler) {
		h.importIntoJobLimit = limit
	}
}

func WithBatchWorkerCount(count int) option {
	return func(h *handler) {
		h.batchWorkerCount = count
	}
}

func WithDDLWorkerCount(count int) option {
	return func(h *handler) {
		h.ddlWorkerCount = count
	}
}

type jobEventHandler struct {
	logger logr.Logger
}

func (h *jobEventHandler) OnAdd(obj any, _ bool) {
	job := h.jobFromInformerObject(obj)
	if job == nil {
		return
	}
	h.logger.Info("job added", "namespace", job.Namespace, "name", job.Name, "active", job.Status.Active, "succeeded", job.Status.Succeeded, "failed", job.Status.Failed)
}

func (h *jobEventHandler) OnUpdate(oldObj, newObj any) {
	newJob := h.jobFromInformerObject(newObj)
	oldJob := h.jobFromInformerObject(oldObj)
	if newJob == nil || oldJob == nil {
		return
	}
	if newJob.ResourceVersion == oldJob.ResourceVersion {
		return
	}
	h.logger.Info(
		"job updated",
		"namespace", newJob.Namespace,
		"name", newJob.Name,
		"active", newJob.Status.Active,
		"succeeded", newJob.Status.Succeeded,
		"failed", newJob.Status.Failed,
	)
}

func (h *jobEventHandler) OnDelete(obj any) {
	job := h.jobFromInformerObject(obj)
	if job == nil {
		return
	}
	h.logger.Info("job deleted", "namespace", job.Namespace, "name", job.Name)
}

func (h *jobEventHandler) jobFromInformerObject(obj any) *batchv1.Job {
	switch t := obj.(type) {
	case *batchv1.Job:
		return t
	case cache.DeletedFinalStateUnknown:
		if job, ok := t.Obj.(*batchv1.Job); ok {
			return job
		}
	}
	return nil
}

func WithImportIntoWorkerCount(count int) option {
	return func(h *handler) {
		h.importIntoWorkerCount = count
	}
}

func WithComponentNode(componentNode string) option {
	return func(h *handler) {
		h.componentNode = componentNode
	}
}

func WithEnableWorkerAutoScale(v bool) option {
	return func(h *handler) {
		h.enableWorkerAutoScale = v
	}
}

func WithMaxSharedWorkerCount(count int) option {
	return func(h *handler) {
		h.maxSharedWorkerCount = count
	}
}

func WithMaxSharedWorkerCountPerKeyspace(count int) option {
	return func(h *handler) {
		h.maxSharedWorkerCountPerKeyspace = count
	}
}

func WithDisableBatchWorkerAutoScale(v bool) option {
	return func(h *handler) {
		h.disableWorkerAutoScaleTaskType[string(job.TypeBatch)] = v
	}
}

func WithDisableDDLWorkerAutoScale(v bool) option {
	return func(h *handler) {
		h.disableWorkerAutoScaleTaskType[string(job.TypeDDL)] = v
	}
}

func WithDisableImportIntoWorkerAutoScale(v bool) option {
	return func(h *handler) {
		h.disableWorkerAutoScaleTaskType[string(job.TypeImportInto)] = v
	}
}