package k8s

import (
	"encoding/json"
	"fmt"
	"strings"

	"github.com/tidbcloud/infra-provider/regional-server/pkg/consts"
	"gomodules.xyz/jsonpatch/v2"
	"k8s.io/apimachinery/pkg/types"
	"sigs.k8s.io/controller-runtime/pkg/client"
)

type ObjectPatch = func(client.Object) *jsonpatch.Operation

func SetAnnotationPatch(key, value string) ObjectPatch {
	return func(obj client.Object) *jsonpatch.Operation {
		annotations, op := obj.GetAnnotations(), "add"
		if annotations == nil {
			return &jsonpatch.Operation{
				Operation: "add",
				Path:      "/metadata/annotations",
				Value:     map[string]string{key: value},
			}
		}
		if annotations[consts.AnnotationKeyConfirmDeletion] != "" {
			op = "replace"
		}
		return &jsonpatch.Operation{
			Operation: op,
			Path:      fmt.Sprintf("/metadata/annotations/%s", strings.ReplaceAll(key, "/", "~1")),
			Value:     value,
		}
	}
}

func RemoveAnnotationPatch(key string) ObjectPatch {
	return func(obj client.Object) *jsonpatch.Operation {
		annotations := obj.GetAnnotations()
		if annotations == nil {
			return nil
		}
		if _, ok := annotations[key]; !ok {
			return nil
		}
		return &jsonpatch.Operation{
			Operation: "remove",
			Path:      fmt.Sprintf("/metadata/annotations/%s", strings.ReplaceAll(key, "/", "~1")),
		}
	}
}

func ConstructPatch(patchPtrs ...*jsonpatch.Operation) (client.Patch, error) {
	patches := make([]jsonpatch.Operation, 0, len(patchPtrs))
	for _, patchPtr := range patchPtrs {
		if patchPtr != nil {
			patches = append(patches, *patchPtr)
		}
	}
	buf, err := json.Marshal(patches)
	if err != nil {
		return nil, err
	}
	return client.RawPatch(types.JSONPatchType, buf), nil
}

var CordonNode ObjectPatch = func(obj client.Object) *jsonpatch.Operation {
	return &jsonpatch.Operation{
		Operation: "replace",
		Path:      "/spec/unschedulable",
		Value:     true,
	}
}

var UncordonNode ObjectPatch = func(obj client.Object) *jsonpatch.Operation {
	return &jsonpatch.Operation{
		Operation: "replace",
		Path:      "/spec/unschedulable",
		Value:     false,
	}
}