package gc

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/gc"
	"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"
)

var _ types.Handler = &handler{}

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

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(&gc.Task{}, &gc.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)
	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")
			}
		}
	}
}

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

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

	// Calculate how many jobs we can create.
	remainingJobsQuota, err := job.CheckJobQuota(
		h.clusterCli, h.namespace, types2.JobTypeLabelKey, string(job.TypeGC), 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 a global gc safe point from PD.
	sp, err := h.pdCli.UpdateGCSafePoint(ctx, 0)
	if err != nil {
		return err
	}

	// Acquire pending tasks.
	var pendingTasks []struct {
		KeyspaceName string
		TidbPool     string
		RegisteredAt time.Time
	}
	err = db.
		Model(&gc.Task{}).
		Select("keyspace_name, tidb_pool, MIN(registered_at) as registered_at").
		Where("state = ?", gc.StatePending).
		Where("deletion_ts < ?", sp).
		Group("keyspace_name").
		Group("tidb_pool").
		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
		}

		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
		}

		go func(keyspaceName string, tidbPool string, registeredAt time.Time) {
			h.logger.Info("activating job", "keyspace-name", keyspaceName)
			err = h.activateJob(ctx, keyspaceName, meta.Id, tidbPool, sp)
			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.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(&gc.Task{}).
			Where("keyspace_name = ?", k).
			Update("state", gc.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 []gc.Task
	)

	// Collect the runners that have finished their tasks.
	err := db.Model(&gc.Task{}).
		Where("state = ?", gc.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(&gc.TaskDone{
				KeyspaceName: doneTask.KeyspaceName,
				TidbPool:     doneTask.TidbPool,
				DeletionTs:   doneTask.DeletionTs,
				Runner:       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[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(&gc.Task{}).
		Distinct("keyspace_name, runner").
		Where("state = ?", gc.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(&gc.Task{}).
				Where("keyspace_name = ?", r.KeyspaceName).
				Where("runner = ?", r.Runner).
				Update("state", gc.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) 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 gc.State
		Count int
	}

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

	countMap := map[gc.State]float64{
		gc.StatePending: 0,
		gc.StateRunning: 0,
		gc.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("gc"),
		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
		}
	}
}