/*
Copyright (C) 2022-2025 ApeCloud Co., Ltd

This file is part of KubeBlocks project

This program is free software: you can redistribute it and/or modify
it under the terms of the GNU Affero General Public License as published by
the Free Software Foundation, either version 3 of the License, or
(at your option) any later version.

This program is distributed in the hope that it will be useful
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
GNU Affero General Public License for more details.

You should have received a copy of the GNU Affero General Public License
along with this program.  If not, see <http://www.gnu.org/licenses/>.
*/

package custom

import (
	"bytes"
	"context"
	"fmt"
	"sort"
	"strconv"
	"strings"

	corev1 "k8s.io/api/core/v1"
	"k8s.io/client-go/util/jsonpath"
	"sigs.k8s.io/controller-runtime/pkg/client"

	appsv1 "github.com/apecloud/kubeblocks/apis/apps/v1"
	opsv1alpha1 "github.com/apecloud/kubeblocks/apis/operations/v1alpha1"
	"github.com/apecloud/kubeblocks/pkg/common"
	"github.com/apecloud/kubeblocks/pkg/constant"
	"github.com/apecloud/kubeblocks/pkg/controller/component"
	"github.com/apecloud/kubeblocks/pkg/controller/scheduling"
	intctrlutil "github.com/apecloud/kubeblocks/pkg/controllerutil"
)

const (
	KBOpsActionNameLabelKey = "opsrequest.kubeblocks.io/action-name"

	KBEnvOpsName             = "KB_OPS_NAME"
	KbEnvOpsNamespace        = "KB_OPS_NAMESPACE"
	kbEnvCompHeadlessSVCName = "KB_COMP_HEADLESS_SVC_NAME"
	kbEnvCompSVCName         = "KB_COMP_SVC_NAME"
	kbEnvCompSVCPortPrefix   = "KB_COMP_SVC_PORT_"
	kbEnvCompServiceVersion  = "KB_COMP_SERVICE_VERSION"
	kbEnvAccountUserName     = "KB_ACCOUNT_USERNAME"
	kbEnvAccountPassword     = "KB_ACCOUNT_PASSWORD"
)

func getComponentInfo(opsDef *opsv1alpha1.OpsDefinition, compDefName string) *opsv1alpha1.ComponentInfo {
	if opsDef == nil {
		return nil
	}
	for _, v := range opsDef.Spec.ComponentInfos {
		if component.PrefixOrRegexMatched(compDefName, v.ComponentDefinitionName) {
			return &v
		}
	}
	return nil
}

// buildComponentDefEnvs builds the env vars by the opsDefinition.spec.componentDefinitionRef
func buildComponentEnvs(reqCtx intctrlutil.RequestCtx,
	cli client.Client,
	cluster *appsv1.Cluster,
	opsDef *opsv1alpha1.OpsDefinition,
	env *[]corev1.EnvVar,
	comp *appsv1.ClusterComponentSpec,
	componentName string) error {
	// inject built-in component env
	fullCompName := constant.GenerateClusterComponentName(cluster.Name, componentName)
	*env = append(*env, []corev1.EnvVar{
		{Name: constant.KBEnvClusterName, Value: cluster.Name},
		{Name: constant.KBEnvCompName, Value: componentName},
		{Name: constant.KBEnvClusterCompName, Value: fullCompName},
		{Name: constant.KBEnvCompReplicas, Value: strconv.Itoa(int(comp.Replicas))},
		{Name: kbEnvCompHeadlessSVCName, Value: constant.GenerateDefaultComponentHeadlessServiceName(cluster.Name, componentName)},
	}...)
	if len(opsDef.Spec.ComponentInfos) == 0 {
		return nil
	}
	// get component definition
	_, compDef, err := component.GetCompNCompDefByName(reqCtx.Ctx, cli, cluster.Namespace, fullCompName)
	if err != nil {
		return err
	}
	componentInfo := getComponentInfo(opsDef, compDef.Name)
	if componentInfo == nil {
		return intctrlutil.NewFatalError(fmt.Sprintf(`componentDefinition "%s" is not support for this operations`, compDef.Name))
	}

	*env = append(*env, corev1.EnvVar{Name: kbEnvCompServiceVersion, Value: compDef.Spec.ServiceVersion})

	buildSecretKeyRef := func(secretName, key string) *corev1.EnvVarSource {
		return &corev1.EnvVarSource{
			SecretKeyRef: &corev1.SecretKeySelector{
				LocalObjectReference: corev1.LocalObjectReference{
					Name: secretName,
				},
				Key: key,
			},
		}
	}
	// inject connect envs
	if componentInfo.AccountName != "" {
		accountSecretName := constant.GenerateAccountSecretName(cluster.Name, componentName, componentInfo.AccountName)
		*env = append(*env, corev1.EnvVar{Name: kbEnvAccountUserName, Value: componentInfo.AccountName})
		*env = append(*env, corev1.EnvVar{Name: kbEnvAccountPassword, ValueFrom: buildSecretKeyRef(accountSecretName, constant.AccountPasswdForSecret)})
	}

	// inject SVC and SVC ports
	if componentInfo.ServiceName != "" {
		for _, v := range compDef.Spec.Services {
			if v.Name != componentInfo.ServiceName {
				continue
			}
			*env = append(*env, corev1.EnvVar{Name: kbEnvCompSVCName, Value: constant.GenerateComponentServiceName(cluster.Name, componentName, v.ServiceName)})
			for _, port := range v.Spec.Ports {
				portName := strings.ReplaceAll(port.Name, "-", "_")
				*env = append(*env, corev1.EnvVar{Name: kbEnvCompSVCPortPrefix + strings.ToUpper(portName), Value: strconv.Itoa(int(port.Port))})
			}
			break
		}
	}
	return nil
}

// BuildEnvVars builds the env vars by the vars of the targetPodTemplate.
func buildEnvVars(reqCtx intctrlutil.RequestCtx,
	cli client.Client,
	targetPod *corev1.Pod,
	vars []opsv1alpha1.OpsEnvVar) ([]corev1.EnvVar, error) {
	if len(vars) == 0 {
		return []corev1.EnvVar{}, nil
	}
	// build vars
	var (
		envVars    []corev1.EnvVar
		err        error
		envFromMap = map[string]map[string]*corev1.EnvVar{}
	)
	buildEnvVarByEnvRef := func(envVarRef *opsv1alpha1.EnvVarRef) (*corev1.EnvVar, error) {
		container := intctrlutil.GetPodContainer(targetPod, envVarRef.TargetContainerName)
		if container == nil {
			return nil, intctrlutil.NewFatalError(fmt.Sprintf(`can not find container "%s" of the pod "%s"`, envVarRef.TargetContainerName, targetPod.Name))
		}
		envVar := buildVarWithEnv(targetPod, container, envVarRef.EnvName)
		if envVar == nil {
			// if the var not found in container.env, try to find from container.envFrom.
			envMap := envFromMap[container.Name]
			if envMap != nil {
				return envMap[envVarRef.EnvName], nil
			}
			envMap, err = getEnvVarsFromEnvFrom(reqCtx.Ctx, cli, targetPod.Namespace, container)
			if err != nil {
				return nil, err
			}
			envFromMap[container.Name] = envMap
			envVar = envMap[envVarRef.EnvName]
		}
		return envVar, nil
	}

	for i := range vars {
		var envVar *corev1.EnvVar
		envVarRef := vars[i].ValueFrom.EnvVarRef
		if envVarRef != nil {
			envVar, err = buildEnvVarByEnvRef(envVarRef)
		} else if vars[i].ValueFrom.FieldRef != nil {
			envVar, err = buildVarWithFieldPath(targetPod, vars[i].ValueFrom.FieldRef.FieldPath)
		}
		if envVar == nil {
			if vars[i].Optional != nil && *vars[i].Optional {
				continue
			}
			return nil, intctrlutil.NewFatalError(fmt.Sprintf(`can not find the env "%s" in the container "%s"`, envVarRef.EnvName, envVarRef.TargetContainerName))
		}
		envVar.Name = vars[i].Name
		envVars = append(envVars, *envVar)
	}
	return envVars, nil
}

// BuildVarWithFieldPath builds the env var by field jsonpath of the pod.
func buildVarWithFieldPath(targetPod *corev1.Pod, fieldPath string) (*corev1.EnvVar, error) {
	path := jsonpath.New("jsonpath")
	if err := path.Parse(fmt.Sprintf("{%s}", fieldPath)); err != nil {
		return nil, fmt.Errorf("failed to parse fieldPath %s", fieldPath)
	}
	buff := bytes.NewBuffer([]byte{})
	if err := path.Execute(buff, targetPod); err != nil {
		return nil, fmt.Errorf("failed to execute fieldPath %s", fieldPath)
	}
	return &corev1.EnvVar{
		Value: buff.String(),
	}, nil
}

// GetEnvVarsFromEnvFrom gets the env var by the container envFrom.
func getEnvVarsFromEnvFrom(ctx context.Context, cli client.Client, podNamespace string, container *corev1.Container) (map[string]*corev1.EnvVar, error) {
	envMap := map[string]*corev1.EnvVar{}
	for _, env := range container.EnvFrom {
		prefix := env.Prefix
		if env.ConfigMapRef != nil {
			configMap := &corev1.ConfigMap{}
			if err := cli.Get(ctx, client.ObjectKey{Name: env.ConfigMapRef.Name, Namespace: podNamespace}, configMap); err != nil {
				return nil, err
			}
			for k, v := range configMap.Data {
				name := prefix + k
				envMap[name] = &corev1.EnvVar{Name: name, Value: v}
			}
		} else if env.SecretRef != nil {
			secret := &corev1.Secret{}
			if err := cli.Get(ctx, client.ObjectKey{Name: env.SecretRef.Name, Namespace: podNamespace}, secret); err != nil {
				return nil, err
			}
			for k := range secret.Data {
				name := prefix + k
				envMap[name] = &corev1.EnvVar{Name: name, ValueFrom: &corev1.EnvVarSource{SecretKeyRef: &corev1.SecretKeySelector{
					LocalObjectReference: corev1.LocalObjectReference{Name: env.SecretRef.Name},
					Key:                  k,
				}}}
			}
		}
	}
	return envMap, nil
}

// BuildVarWithEnv builds the env var by the container env.
func buildVarWithEnv(targetPod *corev1.Pod, container *corev1.Container, envName string) *corev1.EnvVar {
	for i := range container.Env {
		env := container.Env[i]
		if env.Name != envName {
			continue
		}
		if env.ValueFrom != nil && env.ValueFrom.FieldRef != nil {
			// handle fieldRef
			value, _ := common.GetFieldRef(targetPod, env.ValueFrom)
			return &corev1.EnvVar{Name: envName, Value: value}
		}
		return &env
	}
	return nil
}

func buildActionPodEnv(reqCtx intctrlutil.RequestCtx,
	cli client.Client,
	cluster *appsv1.Cluster,
	opsDef *opsv1alpha1.OpsDefinition,
	ops *opsv1alpha1.OpsRequest,
	comp *appsv1.ClusterComponentSpec,
	compCustomItem *opsv1alpha1.CustomOpsComponent,
	podInfoExtractor *opsv1alpha1.PodInfoExtractor,
	targetPod *corev1.Pod) ([]corev1.EnvVar, error) {
	var env = []corev1.EnvVar{
		{
			Name:  KBEnvOpsName,
			Value: ops.Name,
		},
		{
			Name:  KbEnvOpsNamespace,
			Value: ops.Namespace,
		},
	}

	componentName := compCustomItem.ComponentName
	if targetPod != nil {
		// "component" might be a shard component and the name is logical.
		// we should get the real component name from the pod labels.
		componentName = targetPod.Labels[constant.KBAppComponentLabelKey]
	}

	// inject component and componentDef envs
	if err := buildComponentEnvs(reqCtx, cli, cluster, opsDef, &env, comp, componentName); err != nil {
		return nil, err
	}

	if podInfoExtractor != nil {
		// inject vars
		envVars, err := buildEnvVars(reqCtx, cli, targetPod, podInfoExtractor.Env)
		if err != nil {
			return nil, err
		}
		env = append(env, envVars...)
	}

	// inject params env
	for i := range compCustomItem.Parameters {
		param := compCustomItem.Parameters[i]
		if param.Value != "" || param.ValueFrom == nil {
			env = append(env, corev1.EnvVar{Name: param.Name, Value: param.Value})
			continue
		}
		switch {
		case param.ValueFrom.ConfigMapKeyRef != nil:
			env = append(env, corev1.EnvVar{Name: param.Name, ValueFrom: &corev1.EnvVarSource{
				ConfigMapKeyRef: param.ValueFrom.ConfigMapKeyRef,
			}})
		case param.ValueFrom.SecretKeyRef != nil:
			env = append(env, corev1.EnvVar{Name: param.Name, ValueFrom: &corev1.EnvVarSource{
				SecretKeyRef: param.ValueFrom.SecretKeyRef,
			}})
		}
	}
	return env, nil
}

func buildActionPodName(opsRequest *opsv1alpha1.OpsRequest,
	compName,
	actionName string,
	index,
	retries int) string {
	return fmt.Sprintf("%s-%s-%s-%s-%d-%d", opsRequest.UID[:8],
		common.CutString(opsRequest.Name, 30), compName, actionName, index, retries)
}

// getTargetPodInfoExtractor gets the target pod template.
func getTargetPodInfoExtractor(opsDef *opsv1alpha1.OpsDefinition, podInfoExtractorName string) *opsv1alpha1.PodInfoExtractor {
	for _, v := range opsDef.Spec.PodInfoExtractors {
		if podInfoExtractorName == v.Name {
			return &v
		}
	}
	return nil
}

func getTargetTemplateAndPod(ctx context.Context,
	cli client.Client,
	opsDef *opsv1alpha1.OpsDefinition,
	podInfoExtractorName,
	targetPodName,
	podNamespace string) (*opsv1alpha1.PodInfoExtractor, *corev1.Pod, error) {
	if targetPodName == "" {
		return nil, nil, nil
	}
	targetPodTemplate := getTargetPodInfoExtractor(opsDef, podInfoExtractorName)
	targetPod := &corev1.Pod{}
	if err := cli.Get(ctx,
		client.ObjectKey{Name: targetPodName, Namespace: podNamespace}, targetPod); err != nil {
		return nil, nil, err
	}
	return targetPodTemplate, targetPod, nil
}

// getTargetPods gets the target pods by podSelector of the targetPodTemplate.
func getTargetPods(
	ctx context.Context,
	cli client.Client,
	cluster *appsv1.Cluster,
	podSelector opsv1alpha1.PodSelector,
	compName string) ([]*corev1.Pod, error) {
	var (
		pods []*corev1.Pod
		err  error
	)
	if cluster.Spec.GetShardingByName(compName) != nil {
		// get pods of the sharding components
		podList := &corev1.PodList{}
		labels := constant.GetClusterLabels(cluster.Name)
		labels[constant.KBAppShardingNameLabelKey] = compName
		if podSelector.Role != "" {
			labels[constant.RoleLabelKey] = podSelector.Role
		}
		if err = cli.List(ctx, podList, client.InNamespace(cluster.Namespace), client.MatchingLabels(labels)); err != nil {
			return nil, err
		}
		for i := range podList.Items {
			pods = append(pods, &podList.Items[i])
		}
	} else {
		if podSelector.Role != "" {
			pods, err = component.ListOwnedPodsWithRole(ctx, cli, cluster.Namespace, cluster.Name, compName, podSelector.Role)
		} else {
			pods, err = component.ListOwnedPods(ctx, cli, cluster.Namespace, cluster.Name, compName)
		}
	}
	if err != nil {
		return nil, err
	}
	if len(pods) == 0 {
		return nil, intctrlutil.NewFatalError("can not find any pod which matches the podSelector for the component " + compName)
	}
	sort.Sort(intctrlutil.ByPodName(func() []corev1.Pod {
		l := make([]corev1.Pod, 0)
		for i := range pods {
			l = append(l, *pods[i])
		}
		return l
	}()))
	var targetPods []*corev1.Pod
	for i := range pods {
		pod := pods[i]
		targetPods = append(targetPods, pod)
		// Preferably select available pod.
		if podSelector.MultiPodSelectionPolicy == opsv1alpha1.Any && intctrlutil.IsPodAvailable(pod, 0) {
			return []*corev1.Pod{pod}, nil
		}
	}
	if podSelector.MultiPodSelectionPolicy == opsv1alpha1.Any {
		return targetPods[0:1], nil
	}
	return targetPods, nil
}

func buildLabels(opsName, actionName string) map[string]string {
	return map[string]string{
		constant.OpsRequestNameLabelKey: opsName,
		KBOpsActionNameLabelKey:         actionName,
		constant.AppManagedByLabelKey:   constant.AppName,
	}
}

func getNameFromObjectKey(objectKey string) string {
	strs := strings.Split(objectKey, "/")
	if len(strs) == 2 {
		return strs[1]
	}
	return objectKey
}

func getTolerations(cluster *appsv1.Cluster, compSpec *appsv1.ClusterComponentSpec) ([]corev1.Toleration, error) {
	schedulePolicy, err := scheduling.BuildSchedulingPolicy(cluster, compSpec)
	if err != nil {
		return nil, err
	}
	if schedulePolicy == nil {
		return nil, nil
	}
	return schedulePolicy.Tolerations, nil
}
