package main

import (
	"context"
	"flag"
	"os"
	"time"

	corev1 "k8s.io/api/core/v1"
	storagev1 "k8s.io/api/storage/v1"
	"sigs.k8s.io/controller-runtime/pkg/builder"
	"sigs.k8s.io/controller-runtime/pkg/cache"
	"sigs.k8s.io/controller-runtime/pkg/client"
	"sigs.k8s.io/controller-runtime/pkg/client/config"
	"sigs.k8s.io/controller-runtime/pkg/controller"
	"sigs.k8s.io/controller-runtime/pkg/event"
	"sigs.k8s.io/controller-runtime/pkg/healthz"
	logf "sigs.k8s.io/controller-runtime/pkg/log"
	"sigs.k8s.io/controller-runtime/pkg/log/zap"
	"sigs.k8s.io/controller-runtime/pkg/manager"
	"sigs.k8s.io/controller-runtime/pkg/manager/signals"
	metricsserver "sigs.k8s.io/controller-runtime/pkg/metrics/server"
	"sigs.k8s.io/controller-runtime/pkg/predicate"
	"sigs.k8s.io/controller-runtime/pkg/reconcile"
)

const (
	requeueWaitRegister = 2 * time.Second
)

var (
	syncPeriod time.Duration
	concurrent int

	notReadyTaint = &corev1.Taint{
		Key:    "csi.tidbcloud.com/not-ready",
		Value:  "true",
		Effect: corev1.TaintEffectNoSchedule,
	}
)

func main() {
	flag.DurationVar(&syncPeriod, "sync-period", 0,
		"Determines the minimum frequency at which watched resources are reconciled")
	flag.IntVar(&concurrent, "concurrent", 3,
		"Maximum number of concurrent Reconciles")
	opt := zap.Options{}
	opt.BindFlags(flag.CommandLine)
	flag.Parse()

	logf.SetLogger(zap.New(zap.UseFlagOptions(&opt)))
	var log = logf.Log.WithName("csi-node-taint")

	mgr, err := manager.New(config.GetConfigOrDie(), manager.Options{
		LeaderElection:          true,
		LeaderElectionID:        "csi-node-taint",
		LeaderElectionNamespace: "kube-system",
		HealthProbeBindAddress:  ":6001",
		Metrics:                 metricsserver.Options{BindAddress: "0"}, // disable metrics
		Cache:                   cache.Options{SyncPeriod: &syncPeriod},
		Client: client.Options{
			Cache: &client.CacheOptions{DisableFor: []client.Object{&corev1.Node{}}},
		},
	})
	if err != nil {
		log.Error(err, "could not create manager")
		os.Exit(1)
	}
	if err := mgr.AddHealthzCheck("healthz", healthz.Ping); err != nil {
		log.Error(err, "unable to set up healthz check")
		os.Exit(1)
	}
	if err := mgr.AddReadyzCheck("readyz", healthz.Ping); err != nil {
		log.Error(err, "unable to set up readyz check")
		os.Exit(1)
	}

	err = builder.
		ControllerManagedBy(mgr). // Create the ControllerManagedBy
		For(&storagev1.CSINode{}).
		WithEventFilter(&predicate.Funcs{ // only watch Create
			UpdateFunc:  func(updateEvent event.UpdateEvent) bool { return false },
			DeleteFunc:  func(deleteEvent event.DeleteEvent) bool { return false },
			GenericFunc: func(genericEvent event.GenericEvent) bool { return false },
		}).
		WithOptions(controller.Options{MaxConcurrentReconciles: concurrent}).
		Complete(&CSINodeReconciler{})
	if err != nil {
		log.Error(err, "could not create controller")
		os.Exit(1)
	}

	if err := mgr.Start(signals.SetupSignalHandler()); err != nil {
		log.Error(err, "could not start manager")
		os.Exit(1)
	}
}

type CSINodeReconciler struct {
	client.Client
}

func (r *CSINodeReconciler) Reconcile(ctx context.Context, req reconcile.Request) (reconcile.Result, error) {
	log := logf.FromContext(ctx)
	log.Info("start reconcile")
	start := time.Now()
	defer func() {
		log.Info("finish reconcile", "time", time.Since(start))
	}()

	var err error

	csiNode := &storagev1.CSINode{}
	if err = r.Get(ctx, req.NamespacedName, csiNode); err != nil {
		return reconcile.Result{}, err
	}
	if csiNode.Spec.Drivers == nil {
		log.Info("wait csi drivers to register...")
		return reconcile.Result{RequeueAfter: requeueWaitRegister}, nil
	}

	node := &corev1.Node{}
	if err = r.Get(ctx, req.NamespacedName, node); err != nil {
		return reconcile.Result{}, err
	}
	if newTaints, deleted := DeleteTaint(node.Spec.Taints, notReadyTaint); deleted {
		log.Info("delete notReadyTaint")
		node.Spec.Taints = newTaints
		err = r.Update(ctx, node)
	}

	return reconcile.Result{}, err
}

func (r *CSINodeReconciler) InjectClient(c client.Client) error { //nolint:unparam // by design
	r.Client = c
	return nil
}

// DeleteTaint removes all the taints that have the same key and effect to given taintToDelete.
// Copy from pkg/util/taints/taints.go
func DeleteTaint(taints []corev1.Taint, taintToDelete *corev1.Taint) ([]corev1.Taint, bool) {
	newTaints := []corev1.Taint{}
	deleted := false
	for i := range taints {
		if taintToDelete.MatchTaint(&taints[i]) {
			deleted = true
			continue
		}
		newTaints = append(newTaints, taints[i])
	}
	return newTaints, deleted
}