package gcv2

import (
	"context"
	"sync"
	"time"

	"github.com/go-logr/logr"
	"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/resourcemanager/scheduler"
	"github.com/tidbcloud/aws-shared-provider/pkg/tidbworker/schema/gcv2"
	"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"
	"sigs.k8s.io/controller-runtime/pkg/client"
	ctrl "sigs.k8s.io/controller-runtime/pkg/manager"
)

const (
	pollInterval                = 15 * time.Second
	recycleDoneJobsInterval     = 30 * time.Second
	recycleOrphanJobsInterval   = 5 * time.Minute
	recycleJammedWorkerInterval = 1 * time.Minute
	maxWorkerLifetime           = 15 * time.Minute
	requeueDelay                = 15 * time.Minute
	updateMetricsInterval       = 15 * time.Second

	gcOffsetSecond = 300
)

type handler struct {
	clusterCli       client.Client
	namespace        string
	logger           logr.Logger
	wg               sync.WaitGroup
	db               *gorm.DB
	pdCli            pd.Client
	jobBuilder       *job.Builder
	adminConfig      *admin.ConfigReconciler
	workerCertSecret string
	jobLimit         int
	taskRM           *scheduler.Scheduler
}

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(&gcv2.Task{}, &gcv2.TaskDone{})
	if err != nil {
		return err
	}

	initMetrics()

	h.wg.Add(5) //nolint:mnd // number of loops
	go h.pollLoop(ctx)
	go h.recycleDoneJobsLoop(ctx)
	go h.recycleOrphanJobsLoop(ctx)
	go h.recycleJammedWorkerLoop(ctx)
	go h.updateMetricsLoop(ctx)
	h.wg.Wait()

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

func (h *handler) poll(ctx context.Context) error {
	db := h.db.WithContext(ctx)

	if h.adminConfig != nil && h.adminConfig.IsPaused("gcv2") {
		h.logger.Info("skip poll tasks", "handler", "gcv2")
		return nil
	}

	// Calculate how many jobs we can create.
	remainingJobsQuota, err := job.CheckJobQuota(
		h.clusterCli, h.namespace, types2.JobTypeLabelKey, string(job.TypeGCV2), h.jobLimit)
	if err != nil {
		h.logger.Error(err, "failed to count running jobs")
		return err
	}
	if remainingJobsQuota <= 0 {
		h.logger.Info("job limit reached, skip poll loop", "job-limit", h.jobLimit)
		return nil
	}

	// Acquire pending tasks.
	var pendingTasks []struct {
		KeyspaceName  string
		LastGcRunTime int64
		TidbPool      string
		RegisteredAt  time.Time
	}
	err = db.Raw(`
SELECT t0.keyspace_name,
       t0.tidb_pool, 
       t0.last_gc_run_time,
       t0.registered_at
FROM gc_v2_tasks AS t0
JOIN (
    SELECT keyspace_name, 
           MAX(last_gc_run_time) AS max_last_gc_run_time
    FROM gc_v2_tasks
    WHERE state = ?
    GROUP BY keyspace_name
) AS t1 ON t0.keyspace_name = t1.keyspace_name 
        AND t0.last_gc_run_time = t1.max_last_gc_run_time
WHERE t0.state = ?
AND t0.last_gc_run_time < UNIX_TIMESTAMP() - (t0.gc_life_time + ?)
ORDER BY t0.registered_at ASC
`, gcv2.StatePending, gcv2.StatePending, gcOffsetSecond).Scan(&pendingTasks).Error
	if err != nil {
		h.logger.Error(err, "failed to query pendingTasks")
		return err
	}

	disabled := make(map[string]struct{})
	for _, task := range pendingTasks {
		if _, ok := disabled[task.KeyspaceName]; ok {
			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 worker, keyspace is not enabled", "keyspace-name", task.KeyspaceName)
			disabled[task.KeyspaceName] = struct{}{}
			continue
		}

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

		remainingJobsQuota -= 1
		if remainingJobsQuota <= 0 {
			h.logger.Info("job limit reached, skip poll loop", "job-limit", h.jobLimit)
			break
		}
	}
	for k := range disabled {
		err := db.Model(&gcv2.Task{}).
			Where("keyspace_name = ?", k).
			Update("state", gcv2.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")
			}
		}
	}
}

func (h *handler) recycleDoneJobs(ctx context.Context) error {
	var (
		db        = h.db.WithContext(ctx)
		doneTasks []gcv2.Task
	)

	// Collect the runners that have finished their tasks.
	err := db.Model(&gcv2.Task{}).
		Where("state = ?", gcv2.StateDone).
		Find(&doneTasks).Error
	if err != nil {
		h.logger.Error(err, "failed to list finished runner")
		return err
	}

	// deletedRunners is used to mark runners that have been terminated.
	deletedRunners := map[string]struct{}{"": {}}
	for _, doneTask := range doneTasks {
		// Delete target runners.
		// If the job is terminated, ignore this record.
		runner := doneTask.Runner
		if _, ok := deletedRunners[runner]; !ok {
			h.logger.Info("deleting job", "job", runner)
			err = h.recycleJob(ctx, doneTask.KeyspaceName, runner)
			if client.IgnoreNotFound(err) != nil {
				h.logger.Error(err, "failed to delete job", "job", runner)
				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(&gcv2.TaskDone{
				KeyspaceName:  doneTask.KeyspaceName,
				TidbPool:      doneTask.TidbPool,
				LastGCRunTime: doneTask.LastGCRunTime,
				LastGCRunTs:   doneTask.LastGCRunTs,
				Runner:        runner,
				RunnerCreated: doneTask.RunnerCreated,
				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[runner] = struct{}{}
		taskCompleteDuration.Observe(time.Since(doneTask.RegisteredAt).Seconds())
	}
	return nil
}

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(&gcv2.Task{}).
		Distinct("keyspace_name, runner").
		Where("state = ?", gcv2.StateRunning).
		Scan(&results).Error
	if err != nil {
		h.logger.Error(err, "failed to list 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(&gcv2.Task{}).
				Where("keyspace_name = ?", r.KeyspaceName).
				Where("runner = ?", r.Runner).
				Update("state", gcv2.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) recycleJammedWorkerLoop(ctx context.Context) {
	defer h.wg.Done()
	ticker := time.NewTicker(recycleJammedWorkerInterval)
	defer ticker.Stop()

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

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

	err := db.Transaction(func(tx *gorm.DB) error {
		var (
			runners []struct {
				KeyspaceName string
				TidbPool     string
				Runner       string
			}
			now       = time.Now()
			requeueTs = now.Add(requeueDelay)
		)

		err := tx.Model(&gcv2.Task{}).
			Select("keyspace_name", "tidb_pool", "runner").
			Where("state = ?", gcv2.StateRunning).
			Where("runner_created < ?", now.Add(-maxWorkerLifetime).Unix()).
			Scan(&runners).
			Error
		if err != nil {
			h.logger.Error(err, "failed to list jammed workers")
			return err
		}

		for _, r := range runners {
			err := tx.Model(&gcv2.Task{}).
				Where("runner = ?", r.Runner).
				Update("state", gcv2.StateDone).
				Error

			if err != nil {
				return err
			}

			err = tx.Clauses(clause.OnConflict{
				DoNothing: true,
			}).Create(&gcv2.Task{
				KeyspaceName:  r.KeyspaceName,
				TidbPool:      r.TidbPool,
				LastGCRunTime: requeueTs.Unix(),
				LastGCRunTs:   0,
				State:         gcv2.StatePending,
			}).Error
			if err != nil {
				return err
			}
		}

		return nil
	})

	return err
}

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 gcv2.State
		Count int
	}

	var taskCounts []TaskCount
	err := db.Model(&gcv2.Task{}).
		Select("state, count(DISTINCT keyspace_name) as count").
		Group("state").
		Scan(&taskCounts).Error
	if err != nil {
		h.logger.Error(err, "failed to count tasks")
		return
	}

	countMap := map[gcv2.State]float64{
		gcv2.StatePending: 0,
		gcv2.StateRunning: 0,
		gcv2.StateDone:    0,
	}

	for _, tc := range taskCounts {
		countMap[tc.State] = float64(tc.Count)
	}

	for state, count := range countMap {
		currentTaskCounter[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(),
		namespace:  namespace,
		logger:     mgr.GetLogger().WithName("gcv2"),
		db:         db,
		pdCli:      pdCli,
		jobBuilder: job.NewBuilder(namespace, mgr.GetClient()),
	}
	for _, op := range ops {
		op(h)
	}
	return h
}

type option func(h *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 WithJobLimit(limit int) option {
	return func(h *handler) {
		h.jobLimit = limit
	}
}

func WithResourceManager(taskRM *scheduler.Scheduler, enable bool) option {
	return func(h *handler) {
		if enable {
			h.taskRM = taskRM
		}
	}
}