package core

import (
	"context"
	"encoding/json"
	"flag"
	"fmt"
	"path/filepath"
	"sort"
	"sync"
	"time"

	"github.com/guo-shaoge/tiflash-autoscale/pkg/config"
	"github.com/guo-shaoge/tiflash-autoscale/pkg/util/log/logger"
	kruiseapi "github.com/openkruise/kruise-api/apps/v1alpha1"
	kruiseclientset "github.com/openkruise/kruise-api/client/clientset/versioned"
	"go.uber.org/zap"
	k8sapiv1 "k8s.io/api/core/v1"
	k8sapierrors "k8s.io/apimachinery/pkg/api/errors"
	k8smetav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/labels"
	"k8s.io/apimachinery/pkg/types"
	k8sapiwatch "k8s.io/apimachinery/pkg/watch"
	k8s "k8s.io/client-go/kubernetes"
	"k8s.io/client-go/rest"
	restclient "k8s.io/client-go/rest"
	"k8s.io/client-go/tools/clientcmd"
	"k8s.io/client-go/util/homedir"
)

const (
	namespace    string = "tiflash-autoscale"
	cloneSetName        = "readnode"
)

// REF: https://openkruise.io/docs/developer-manuals/go-client/
type K8sMgr struct {
	ctx context.Context

	// watch pod status(add/del...)
	k8sCli k8s.Interface
	// lableSelector for cloneset.
	labelSelector k8smetav1.LabelSelector
	cloneSetMu    struct {
		sync.Mutex
		cloneSet *kruiseapi.CloneSet
	}
	cloneSetCli        kruiseclientset.Interface
	podResourceVersion string

	autoScalerConfig *config.Config
	logger           *zap.Logger

	podMgr    *PodMgr
	configMgr *ConfigMgr

	addingCntMu struct {
		sync.Mutex
		cnt int32
	}
}

func NewK8sMgr(ctx context.Context, configMgr *ConfigMgr, parentLogger *zap.Logger) *K8sMgr {
	res := &K8sMgr{
		ctx:       ctx,
		logger:    parentLogger.With(zap.String(logger.LogIDMod, "k8sMgr")),
		configMgr: configMgr,
	}

	if logger.InTest {
		return res
	}

	config, err := getK8sConfig()
	if err != nil {
		panic(err.Error())
	}

	res.k8sCli, err = k8s.NewForConfig(config)
	if err != nil {
		panic(err.Error())
	}
	res.tryCreateNamespaceOrDie()

	res.autoScalerConfig = configMgr.GetConfig()
	res.labelSelector = k8smetav1.LabelSelector{MatchLabels: map[string]string{"app": cloneSetName}}

	res.cloneSetCli = kruiseclientset.NewForConfigOrDie(config)
	res.setupCloneSetOrDie()
	res.podResourceVersion, _, err = res.ListPods()
	if err != nil {
		panic(err)
	}
	return res
}

func (k *K8sMgr) startBgWorker(wg *sync.WaitGroup) {
	wg.Add(1)
	go func() {
		k.cloneSetWatcherLoop()
	}()
	go func() {
		k.watchCNImage(*k.configMgr.GetConfig().GetClusterConfig().SupervisorImage)
	}()
}

func (k *K8sMgr) SetPodMgr(podMgr *PodMgr) {
	k.podMgr = podMgr
}

func (k *K8sMgr) AddK8sPods(argNum uint, retryCnt uint32) (err error) {
	if argNum == 0 {
		return nil
	}

	// To avoid add replicas too fast too many times before pod ready.
	num, oriAddingCnt, needAdd := func() (uint, int32, bool) {
		k.addingCntMu.Lock()
		defer k.addingCntMu.Unlock()

		addingCnt := k.addingCntMu.cnt
		if uint(addingCnt) >= argNum {
			return 0, addingCnt, false
		}

		k.addingCntMu.cnt = int32(argNum)
		return argNum - uint(addingCnt), addingCnt, true
	}()

	if !needAdd || num == 0 {
		k.logger.Info("AddK8sPods prepare, no need to add", zap.Any("ori adding cnt", oriAddingCnt),
			zap.Any("arg num", argNum), zap.Any("res num", num))
		return nil
	}

	var iter int
	updater := func() error {
		k.cloneSetMu.Lock()
		defer k.cloneSetMu.Unlock()

		var (
			newCloneSet *kruiseapi.CloneSet
			updateErr   error
			getErr      error
		)

		newCloneSet, getErr = k.cloneSetCli.AppsV1alpha1().CloneSets(namespace).
			Get(k.ctx, cloneSetName, k8smetav1.GetOptions{})
		if getErr != nil {
			return getErr
		}
		k.cloneSetMu.cloneSet = newCloneSet.DeepCopy()

		oldReplicas := k.cloneSetMu.cloneSet.Spec.Replicas
		newReplicasVal := int32(*oldReplicas) + int32(num)
		k.cloneSetMu.cloneSet.Spec.Replicas = &newReplicasVal
		newCloneSet, updateErr = k.cloneSetCli.AppsV1alpha1().CloneSets(namespace).
			Update(k.ctx, k.cloneSetMu.cloneSet, k8smetav1.UpdateOptions{})

		defer func() {
			k.logger.Info("AddK8sPods iter", zap.Any("iter", iter),
				zap.Any("add num", num), zap.Any("oldReplicas", *oldReplicas),
				zap.Any("newReplicas", newReplicasVal),
				zap.Any("updateErr", updateErr), zap.Any("getErr", getErr))
		}()

		if updateErr != nil {
			newCloneSet, getErr = k.cloneSetCli.AppsV1alpha1().CloneSets(namespace).
				Get(k.ctx, cloneSetName, k8smetav1.GetOptions{})
			if getErr != nil {
				return getErr
			}
			k.cloneSetMu.cloneSet = newCloneSet.DeepCopy()
			return updateErr
		}
		k.cloneSetMu.cloneSet = newCloneSet.DeepCopy()
		return nil
	}

	for iter = 0; iter < int(retryCnt); iter++ {
		if err = updater(); err == nil {
			return nil
		}
		time.Sleep(time.Duration(3) * time.Millisecond)
	}
	return err
}

func updateSkipDelReason(podName string, found bool, k8sPod *k8sapiv1.Pod, reason string) string {
	var newReason string
	if len(reason) != 0 {
		newReason = fmt.Sprintf("%s; pod: %s, found: %v", reason, podName, found)
	} else {
		newReason = fmt.Sprintf("pod: %s, found: %v", podName, found)
	}
	if k8sPod != nil {
		if k8sPod.DeletionTimestamp != nil {
			newReason = fmt.Sprintf("%s, k8sPod.DeletionTimestamp is %v", newReason, *k8sPod.DeletionTimestamp)
		}
		if k8sPod.DeletionGracePeriodSeconds != nil {
			newReason = fmt.Sprintf("%s, k8sPod.DeletionGracePeriodSeconds is %v", newReason, *k8sPod.DeletionGracePeriodSeconds)
		}
	} else {
		newReason += ", k8sPod meta is nil"
	}
	return newReason
}

func (k *K8sMgr) DelK8sPods(argPods []string, warmupPoolCap uint, retryCnt uint32, needChangeReplica bool) (err error) {
	if len(argPods) == 0 {
		return nil
	}

	_, k8sPods, err := k.ListPods()
	if err != nil {
		return err
	}
	k8sPodsStr := printK8sPods(k8sPods)

	pods := make([]string, 0, len(argPods))
	var skipDelReason string
	for _, podName := range argPods {
		k8sPod, ok := k8sPods[podName]
		if !ok || k8sPod.DeletionTimestamp != nil || k8sPod.DeletionGracePeriodSeconds != nil {
			skipDelReason = updateSkipDelReason(podName, ok, k8sPod, skipDelReason)
			continue
		}
		pods = append(pods, podName)
	}

	k.logger.Info("DelK8sPods prepare",
		zap.Any("list k8s pods", k8sPodsStr),
		zap.Any("need change replica", needChangeReplica),
		zap.Any("arg pods", argPods), zap.Any("del pods", pods), zap.Any("skip reason", skipDelReason))
	if len(pods) == 0 {
		return nil
	}

	var iter int
	updater := func() error {
		k.cloneSetMu.Lock()
		defer k.cloneSetMu.Unlock()

		var (
			newCloneSet   *kruiseapi.CloneSet
			updateErr     error
			getErr        error
			correctErrStr string
		)

		newCloneSet, getErr = k.cloneSetCli.AppsV1alpha1().CloneSets(namespace).
			Get(k.ctx, cloneSetName, k8smetav1.GetOptions{})
		if getErr != nil {
			return getErr
		}
		k.cloneSetMu.cloneSet = newCloneSet.DeepCopy()

		// skip delete pod which already in PodsToDelete.
		var newDelNum int32
		oriPodsToDelete := fmt.Sprintf("%v", k.cloneSetMu.cloneSet.Spec.ScaleStrategy.PodsToDelete)
		for _, p1 := range pods {
			var found bool
			for _, p2 := range k.cloneSetMu.cloneSet.Spec.ScaleStrategy.PodsToDelete {
				if p1 == p2 {
					found = true
					break
				}
			}
			if !found {
				k.cloneSetMu.cloneSet.Spec.ScaleStrategy.PodsToDelete = append(
					k.cloneSetMu.cloneSet.Spec.ScaleStrategy.PodsToDelete,
					p1)
				newDelNum++
			}
		}
		newPodsToDelete := fmt.Sprintf("%v", k.cloneSetMu.cloneSet.Spec.ScaleStrategy.PodsToDelete)

		oldReplicas := k.cloneSetMu.cloneSet.Spec.Replicas
		newReplicasVal := *oldReplicas
		if needChangeReplica {
			newReplicasVal = *oldReplicas - newDelNum
			if newReplicasVal < 0 {
				correctErrStr = fmt.Sprintf("newReplicasVal(%v) is less than zero, correct to 1", newReplicasVal)
				newReplicasVal = 1
			}
			if newReplicasVal < int32(warmupPoolCap) {
				correctErrStr += fmt.Sprintf("newReplicasVal(%v) is less than warmupPoolCap(%v), correct to warmupPoolCap", newReplicasVal, warmupPoolCap)
				newReplicasVal = int32(warmupPoolCap)
			}
		}

		defer func() {
			k.logger.Info("DelK8sPods iter", zap.Any("iter", iter),
				zap.Any("del pods", pods),
				zap.Any("oriPodsToDelete", oriPodsToDelete),
				zap.Any("newPodsToDelete", newPodsToDelete),
				zap.Any("new del num", newDelNum),
				zap.Any("oldReplicas", oldReplicas),
				zap.Any("newReplicas", newReplicasVal),
				zap.Any("needChangeReplica", needChangeReplica),
				zap.Any("updateErr", updateErr), zap.Any("getErr", getErr), zap.Any("correctErr", correctErrStr))
		}()

		if needChangeReplica {
			k.cloneSetMu.cloneSet.Spec.Replicas = &newReplicasVal
		}

		newCloneSet, updateErr = k.cloneSetCli.AppsV1alpha1().CloneSets(namespace).
			Update(k.ctx, k.cloneSetMu.cloneSet, k8smetav1.UpdateOptions{})
		if updateErr != nil {
			newCloneSet, getErr = k.cloneSetCli.AppsV1alpha1().CloneSets(namespace).
				Get(k.ctx, cloneSetName, k8smetav1.GetOptions{})
			if getErr != nil {
				return getErr
			}
			k.cloneSetMu.cloneSet = newCloneSet.DeepCopy()
			return updateErr
		}
		k.cloneSetMu.cloneSet = newCloneSet.DeepCopy()
		return nil
	}

	for iter = 0; iter < int(retryCnt); iter++ {
		if err = updater(); err == nil {
			return nil
		}
	}
	return err
}

type sortK8sPodT struct {
	name   string
	status string
}

func printK8sPods(k8sPods map[string]*k8sapiv1.Pod) string {
	s := make([]sortK8sPodT, 0, len(k8sPods))
	for podName, pod := range k8sPods {
		s = append(s, sortK8sPodT{
			name:   podName,
			status: string(pod.Status.Phase),
		})
	}
	sort.Slice(s, func(i, j int) bool {
		return s[i].name < s[j].name
	})
	var resStr string
	for _, pod := range s {
		resStr += fmt.Sprintf("%v,%v; ", pod.name, pod.status)
	}
	return resStr
}

func (k *K8sMgr) ListPods() (string, map[string]*k8sapiv1.Pod, error) {
	pods, err := k.k8sCli.CoreV1().Pods(namespace).List(
		k.ctx,
		k8smetav1.ListOptions{LabelSelector: labels.Set(k.labelSelector.MatchLabels).String()})
	if err != nil {
		return "", nil, err
	}

	k8sPodsMap := make(map[string]*k8sapiv1.Pod)
	for _, pod := range pods.Items {
		k8sPodsMap[pod.Name] = &pod
	}
	return pods.ListMeta.ResourceVersion, k8sPodsMap, nil
}

// patchStringValue specifies a patch operation for a string.
type KubePatchStringValue struct {
	Op    string `json:"op"`
	Path  string `json:"path"`
	Value string `json:"value"`
}

func (k *K8sMgr) watchCNImage(lastCNImage string) {
	if len(lastCNImage) == 0 {
		panic("unexpected supervisor image")
	}
	for {
		clusterConfig := k.configMgr.GetConfig().GetClusterConfig()
		newCNImage := *clusterConfig.SupervisorImage
		if len(lastCNImage) != 0 && len(newCNImage) != 0 && lastCNImage != newCNImage {
			k.logger.Info("got new SupervisorImage",
				zap.String("old", lastCNImage),
				zap.String("new", newCNImage))

			// todo call k8s api to get latest cn image.
			lastCNImage = newCNImage

			payload := []KubePatchStringValue{{
				Op:    "replace",
				Path:  "/spec/template/spec/containers/0/image",
				Value: newCNImage,
			}}
			payloadBytes, err := json.Marshal(payload)
			if err != nil {
				k.logger.Error("patch supervisor image failed", zap.Error(err))
			} else {
				// no need to lock, because it's a patch call.
				_, err = k.cloneSetCli.AppsV1alpha1().CloneSets(namespace).Patch(
					context.TODO(), cloneSetName, types.JSONPatchType, payloadBytes, k8smetav1.PatchOptions{})
				if err != nil {
					k.logger.Error("patch supervisor image failed", zap.Error(err))
				}
			}
		}

		interval := *clusterConfig.ConfigCheckIntervalSec
		select {
		case <-time.After(time.Duration(interval) * time.Second):
		case <-k.ctx.Done():
			k.logger.Info("got ctx.Done, exit")
			return
		}
	}
}

func (k *K8sMgr) cloneSetWatcherLoop() {
	logger := k.logger.With(zap.String("bgWorker", "cloneSetWatcherLoop"))
	for {
		watcher, err := k.k8sCli.CoreV1().Pods(namespace).Watch(
			k.ctx,
			k8smetav1.ListOptions{
				LabelSelector:   labels.Set(k.labelSelector.MatchLabels).String(),
				ResourceVersion: k.podResourceVersion,
			})
		if err != nil {
			k.logger.Error("failed to create watcher ch", zap.Error(err))
			time.Sleep(1 * time.Second)
			continue
		}

		ch := watcher.ResultChan()

	innerloop:
		for {
			select {
			case <-k.ctx.Done():
				logger.Info("got ctx.Done(), exit")
				return
			case event, ok := <-ch:
				if !ok {
					logger.Info("watcher channel close")
					time.Sleep(1 * time.Second)
					tmpPodResourceVersion, _, err := k.ListPods()
					if err == nil {
						k.podResourceVersion = tmpPodResourceVersion
					} else {
						logger.Info("ListPods error", zap.Error(err))
					}
					break innerloop
				}
				k.handleWatchEvent(event)
			}
		}
	}
}

func (k *K8sMgr) handleWatchEvent(event k8sapiwatch.Event) {
	pod, ok := event.Object.(*k8sapiv1.Pod)
	if !ok {
		return
	}
	logger := k.logger.With(zap.String("pod", pod.Name))
	switch event.Type {
	case k8sapiwatch.Added:
		logger.Info("got pod add event")
		func() {
			k.addingCntMu.Lock()
			defer k.addingCntMu.Unlock()

			if k.addingCntMu.cnt > 0 {
				k.addingCntMu.cnt--
			}
		}()

		// Manually deleted pod will be running againg using the same name.
		// So call update instead of add.
		k.podMgr.AddPodDesc(pod)
	case k8sapiwatch.Deleted:
		logger.Info("got pod delete event")
		k.podMgr.DelPodDesc(pod)
	case k8sapiwatch.Modified:
		logger.Info("got pod modified event", zap.Any("pod.Status", pod.Status))
		k.podMgr.UpdatePodDesc(pod)
	case k8sapiwatch.Error:
		logger.Info("got pod error event")
		k.podMgr.DelPodDesc(pod)
	case k8sapiwatch.Bookmark:
		logger.Error("got pod bookmark event, ignore it")
	default:
		logger.Error("got unexpected pod event, ignore it")
	}
}

func (k *K8sMgr) tryCreateNamespaceOrDie() {
	_, err := k.k8sCli.CoreV1().Namespaces().Get(k.ctx, namespace, k8smetav1.GetOptions{})
	if err != nil {
		k.logger.Info("creating namespace", zap.String("ns", namespace))
		_, err := k.k8sCli.CoreV1().Namespaces().Create(k.ctx,
			&k8sapiv1.Namespace{
				ObjectMeta: k8smetav1.ObjectMeta{
					Name: namespace,
					Labels: map[string]string{
						"ns": namespace,
					},
				},
			},
			k8smetav1.CreateOptions{})
		if err != nil {
			panic(err)
		}
	}
}

func (k *K8sMgr) setupCloneSetOrDie() {
	cloneSetInterface := k.cloneSetCli.AppsV1alpha1().CloneSets(namespace)
	template := getCloneSetTemplate(cloneSetName, k.autoScalerConfig, k.logger)

	cloneSet, err := cloneSetInterface.Get(k.ctx, cloneSetName, k8smetav1.GetOptions{})
	if err != nil {
		if !k8sapierrors.IsNotFound(err) {
			panic(err)
		}
		k.logger.Info("start to create readnode cloneset")
		if cloneSet, err = cloneSetInterface.Create(k.ctx, template, k8smetav1.CreateOptions{}); err != nil {
			panic(err)
		}
	} else {
		k.logger.Info("cloneSet already exists, update instead create a new one")
		// change cloneSet.Spec except replicas.
		template.Spec.Replicas = cloneSet.Spec.Replicas
		cloneSet.Spec = template.Spec
		if cloneSet, err = cloneSetInterface.Update(k.ctx, cloneSet, k8smetav1.UpdateOptions{}); err != nil {
			panic(err)
		}
	}

	k.cloneSetMu.cloneSet = cloneSet.DeepCopy()
}

func getK8sConfig() (*restclient.Config, error) {
	if config, err := rest.InClusterConfig(); err != nil {
		return outsideConfig()
	} else {
		return config, err
	}
}

func outsideConfig() (*restclient.Config, error) {
	var kubeconfig *string
	if home := homedir.HomeDir(); home != "" {
		kubeconfig = flag.String("kubeconfig", filepath.Join(home, ".kube", "config"), "(optional) absolute path to the kubeconfig file")
	} else {
		kubeconfig = flag.String("kubeconfig", "", "absolute path to the kubeconfig file")
	}
	flag.Parse()

	// use the current context in kubeconfig
	return clientcmd.BuildConfigFromFlags("", *kubeconfig)
}