package assigncachelimit

import (
	"context"
	"fmt"
	"strconv"
	"time"

	"github.com/tidbcloud/aws-shared-provider/pkg/manager"
	"github.com/tidbcloud/aws-shared-provider/pkg/tidbmanager"
	"github.com/tidbcloud/aws-shared-provider/pkg/types"
	ep "github.com/tidbcloud/aws-shared-provider/pkg/utils/endpoint"
	corev1 "k8s.io/api/core/v1"
	apierrors "k8s.io/apimachinery/pkg/api/errors"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/utils/ptr"
	ctrl "sigs.k8s.io/controller-runtime"
	"sigs.k8s.io/controller-runtime/pkg/builder"
	"sigs.k8s.io/controller-runtime/pkg/client"
	"sigs.k8s.io/controller-runtime/pkg/event"
	"sigs.k8s.io/controller-runtime/pkg/log"
	"sigs.k8s.io/controller-runtime/pkg/predicate"
)

const requeueAfter = 2 * time.Second

// Reconciler trims assigned TiDB pods when assign-cache-size-limit is reduced below current usage.
type Reconciler struct {
	Client      client.Client
	TidbManager tidbmanager.TidbManager
	Namespace   string
}

func (r *Reconciler) Reconcile(ctx context.Context, req ctrl.Request) (ctrl.Result, error) {
	logger := log.FromContext(ctx)
	logger.V(2).Info("start reconcile assign cache limit")
	start := time.Now()
	defer func() {
		logger.V(2).Info("finish reconcile assign cache limit", "time", time.Since(start))
	}()

	limit, err := r.assignCacheSizeLimit(ctx, req)
	if err != nil {
		return ctrl.Result{}, err
	}
	if limit <= 0 {
		return ctrl.Result{}, nil
	}

	_, current := r.TidbManager.AssignCacheSize()
	logger.V(1).Info("assign cache size", "current", current, "limit", limit)
	if current <= limit {
		return ctrl.Result{}, nil
	}

	needDelete := current - limit
	if deleted := r.trimAssignedPods(ctx, needDelete); deleted < needDelete {
		logger.Info("assign cache still exceeds limit after one walk, requeue", "needDelete", needDelete, "deleted", deleted)
		return ctrl.Result{RequeueAfter: requeueAfter}, nil
	}

	return ctrl.Result{}, nil
}

func (r *Reconciler) trimAssignedPods(ctx context.Context, needDelete int) (deleted int) {
	walkDeleted := 0
	endpointCandidates := make([]string, 0, needDelete)
	r.TidbManager.WalkActiveTokenFromAssignCache(func(token string, endpoints []string, replicas int) bool {
		if len(endpoints) == 0 {
			return true
		}
		endpointCandidates = append(endpointCandidates, endpoints[0])
		walkDeleted++
		return walkDeleted < needDelete
	})

	for _, endpoint := range endpointCandidates {
		if r.deleteAssignedPod(ctx, endpoint) == nil {
			deleted++
		}
	}

	return deleted
}

func (r *Reconciler) deleteAssignedPod(ctx context.Context, endpoint string) error {
	podName, _ := ep.ParsePodFromEndpoint(endpoint)

	pod := &corev1.Pod{
		ObjectMeta: metav1.ObjectMeta{
			Name:      podName,
			Namespace: r.Namespace,
		},
	}
	if err := r.Client.Delete(ctx, pod, &client.DeleteOptions{GracePeriodSeconds: ptr.To[int64](0)}); err != nil {
		if err = client.IgnoreNotFound(err); err != nil {
			log.FromContext(ctx).Error(err, "failed to delete assigned tidb pod", "endpoint", endpoint)
		}
		return err
	}

	log.FromContext(ctx).Info("[pod lifecycle] delete assigned tidb pod", "endpoint", endpoint)
	return nil
}

func (r *Reconciler) assignCacheSizeLimit(ctx context.Context, req ctrl.Request) (int, error) {
	cm := &corev1.ConfigMap{}
	if err := r.Client.Get(ctx, req.NamespacedName, cm); err != nil {
		if apierrors.IsNotFound(err) {
			log.FromContext(ctx).Info("configmap not found means unlimited")
			return 0, nil
		}
		log.FromContext(ctx).Error(err, "failed to get configmap")
		return 0, err
	}

	assignCacheSizeLimit := cm.Data[types.AssignCacheSizeLimitKey]
	if assignCacheSizeLimit == "" {
		return 0, nil
	}
	limit, err := strconv.Atoi(assignCacheSizeLimit)
	if err != nil || limit < 0 {
		log.FromContext(ctx).Error(err, "unlimit assign cache size", "limit", limit)
		return 0, nil
	}

	return limit, nil
}

func (r *Reconciler) SetupWithManager(mgr manager.Manager) error {
	if err := ctrl.NewControllerManagedBy(mgr).Named("AssignCacheLimit-"+r.Namespace).
		For(&corev1.ConfigMap{}, builder.WithPredicates(r.configMapPredicate())).
		Complete(r); err != nil {
		return fmt.Errorf("can't setup AssignCacheLimit controller: %w", err)
	}

	return nil
}

func (r *Reconciler) configMapPredicate() predicate.Predicate {
	isTarget := func(object client.Object) bool {
		if object == nil {
			return false
		}
		return object.GetNamespace() == r.Namespace && object.GetName() == types.ServerlessPoolConfig
	}

	return predicate.Funcs{
		CreateFunc: func(createEvent event.CreateEvent) bool {
			return isTarget(createEvent.Object)
		},
		UpdateFunc: func(updateEvent event.UpdateEvent) bool {
			return isTarget(updateEvent.ObjectNew)
		},
		DeleteFunc: func(deleteEvent event.DeleteEvent) bool {
			return false
		},
	}
}