// Package edgerelease implements the EdgeRelease Kpt function, based on the
// Kpt delivery pipeline RFC: https://docs.edge-infra.dev/rfc/kpt-delivery-pipeline/.
// This function can be used to update container image digests for GKE component
// manifests as well as GKE Kpt functions.
//
// Current limitations:
//
// - Cannot handle multiple-image components
package edgerelease

import (
	"fmt"
	"strings"

	"sigs.k8s.io/kustomize/kyaml/fn/runtime/runtimeutil"
	"sigs.k8s.io/kustomize/kyaml/kio/kioutil"
	"sigs.k8s.io/kustomize/kyaml/yaml"

	"edge-infra.dev/pkg/edge/component/build"
	"edge-infra.dev/pkg/edge/component/build/image"
	"edge-infra.dev/pkg/edge/constants"
	fnv1alpha1 "edge-infra.dev/pkg/edge/gitops/fn/v1alpha1"
	"edge-infra.dev/pkg/k8s/meta"
)

// containerPaths are the subpaths within the PodSpec that contain ContainerSpecs
var containerPaths = []string{"containers", "initContainers"}

// Run accepts an input array of yaml.RNodes and updates the resources with the
// configured build metadata + image digest if it corresponds to one of the
// provided component or Kpt function names.
//
// Component manifests are identified by checking
// the provided component name against the value of the
// `platform.edge.ncr.com/component` label.
//
// Kpt functions are identified by their apiVersion (e.g., fns.edge.ncr.com/v1alpha1)
// and the image name in their function configuration annotation.  If the image
// name is present in the list of provided Kpt function images, it is updated
// with the provided digest.
//
// TODO: handle multiple-image Components
func (f *EdgeRelease) Run(input []*yaml.RNode) ([]*yaml.RNode, error) {
	if err := f.Spec.Validate(); err != nil {
		return nil, err
	}

	f.init()

	input, err := kioutil.MapMeta(input, f.componentMapper)
	if err != nil {
		return nil, err
	}

	input, err = kioutil.MapMeta(input, f.kptFnMapper)
	if err != nil {
		return nil, err
	}

	return input, nil
}

func (f *EdgeRelease) componentMapper(i *yaml.RNode, m yaml.ResourceMeta) (*yaml.RNode, error) {
	// determine if this resource belongs to a platform component
	cname, ok := m.Labels[constants.PlatformComponent]
	if !ok {
		return i, nil
	}

	// check if the component is one we are updating
	componentImgs, ok := f.componentMap[cname]
	// we only want to update manifests which describe workloads for our component
	if !ok || !meta.IsWorkload(m.Kind) {
		return i, nil
	}

	// add build metadata labels
	if f.Spec.AddBuildLabels {
		if err := f.addBuildLabels(i, m.Labels); err != nil {
			return nil, wrapErr("edgerelease.componentMapper: failed to add build labels", m, err)
		}
	}

	// grab the correct object path for the []ContainerSpec
	podSpecPath, err := getPodSpecPath(m.Kind)
	if err != nil {
		return nil, wrapErr("edgerelease.componentMapper: failed to determine PodSpec path from workload", m, err)
	}

	// Single anonymous image components cannot support initContainers, so
	// just run the "containers" path and return
	if isDefaultImg(componentImgs) {
		if err := patchContainers(i, componentImgs, append(podSpecPath, "containers")); err != nil {
			return nil, wrapErr("edgerelease.componentMapper: failed to patch containers", m, err)
		}
		return i, nil
	}

	for _, p := range containerPaths {
		if err := patchContainers(i, componentImgs, append(podSpecPath, p)); err != nil {
			return nil, wrapErr("edgerelease.componentMapper: failed to patch containers", m, err)
		}
	}

	return i, nil
}

func patchContainers(i *yaml.RNode, imgs map[string]string, path []string) error {
	// get list of containers
	containerList, err := i.Pipe(yaml.Lookup(path...))
	if err != nil {
		return fmt.Errorf("failed to get list of containers from workload: %w", err)
	}

	// turn it into an array so we can work with it more easily
	containers, err := containerList.Elements()
	if err != nil {
		return fmt.Errorf("failed to get list of containers from workload: %w", err)
	}

	// if there is only one container and we only have one image, we can
	// simply replace the reference.  this is kept here for backwards compatibility
	// with single image workloads that might not provide a container name
	// it should probably be removed long-term once users of EdgeRelease with
	// single container workloads have updated their usage
	// only match this if the default image name is detected
	if isDefaultImg(imgs) && len(containers) == 1 {
		newImg := yaml.NewStringRNode(imgs[defaultImgName])
		return containers[0].PipeE(yaml.SetField("image", newImg))
	}

	for i := range containers {
		c := containers[i]
		name, err := c.Pipe(yaml.Lookup("name"))
		if err != nil {
			return err
		}

		if ref, ok := imgs[strings.TrimSpace(name.MustString())]; ok {
			if err := c.PipeE(yaml.SetField("image", yaml.NewStringRNode(ref))); err != nil {
				return err
			}
		}
	}

	return nil
}

func (f *EdgeRelease) kptFnMapper(i *yaml.RNode, m yaml.ResourceMeta) (*yaml.RNode, error) {
	// check if its an GKE Kpt function
	if m.APIVersion != fnv1alpha1.GroupVersion.String() {
		return i, nil
	}

	fnAnno, ok := m.Annotations[runtimeutil.FunctionAnnotationKey]
	// no function config, we can't update anything
	if !ok {
		return i, nil
	}

	fnConfig, err := yaml.Parse(fnAnno)
	if err != nil {
		return nil, wrapErr("edgerelease.kptFnMapper: failed to parse function annotation", m, err)
	}

	imgNode, err := fnConfig.Pipe(yaml.Lookup("container", "image"))
	if err != nil {
		return nil, wrapErr("edgerelease.kptFnMapper: failed to look up image in function annotation", m, err)
	}

	// if there was nothing found at `container.image` in the Kpt function anno,
	// this function doesn't rely on docker runtime and we can skip it
	if imgNode == nil {
		return i, nil
	}

	img, err := imgNode.String()
	if err != nil {
		return nil, wrapErr("edgerelease.kptFnMapper: failed to get string for function image", m, err)
	}

	kptFn, ok := f.kptFnMap[image.NameFromRef(img)]
	if !ok {
		return i, nil
	}

	if f.Spec.AddBuildLabels {
		if err := f.addBuildLabels(i, m.Labels); err != nil {
			return nil, wrapErr("edgerelease.kptFnMapper: failed to add build labels", m, err)
		}
	}

	newImg := map[string]string{"image": kptFn.Reference()}
	if err := fnConfig.PipeE(yaml.SetField("container", yaml.NewMapRNode(&newImg))); err != nil {
		return nil, wrapErr("edgerelease.kptFnMapper: failed to update function annotation", m, err)
	}

	fnConfigStr, err := fnConfig.String()
	if err != nil {
		return nil, wrapErr("edgerelease.kptFnMapper: failed to convert new function annotation to string", m, err)
	}

	if err := i.PipeE(yaml.SetAnnotation(runtimeutil.FunctionAnnotationKey, fnConfigStr)); err != nil {
		return nil, wrapErr("edgerelease.kptFnMapper: failed to set function annotation", m, err)
	}

	return i, nil
}

// addBuildLabels adds the build metadata from an EdgeRelease struct to an existing
// label map
func (f *EdgeRelease) addBuildLabels(i *yaml.RNode, labels map[string]string) error {
	if labels == nil {
		labels = map[string]string{}
	}
	labels[build.CommitLabel] = f.Spec.Metadata.Commit
	labels[build.SemVerLabel] = f.Spec.Metadata.SemVer
	labels[build.TimestampLabel] = f.Spec.Metadata.Timestamp
	labels[build.BuildIDLabel] = f.Spec.Metadata.ID
	labels[build.RepoLabel] = f.Spec.Metadata.Repo
	labels[build.OrgLabel] = f.Spec.Metadata.Org

	return i.SetLabels(labels)
}

func getPodSpecPath(kind string) ([]string, error) {
	switch kind {
	// all workload types other than CronJobs have the same  path to the pod spec
	case "Deployment", "DaemonSet", "StatefulSet", "Job":
		return []string{"spec", "template", "spec"}, nil
	case "CronJob":
		return []string{"spec", "jobTemplate", "spec", "template", "spec"}, nil
	default:
		return []string{}, fmt.Errorf("could not match %s when looking up PodSpec path", kind)
	}
}

func wrapErr(msg string, m yaml.ResourceMeta, err error) error {
	if path, ok := m.Annotations[kioutil.PathAnnotation]; ok {
		return fmt.Errorf("%s: %s/%s from %s. error: %w", msg, m.Kind, m.Name, path, err)
	}
	return fmt.Errorf("%s: %s/%s. error: %w", msg, m.Kind, m.Name, err)
}

// isDefaultImg detects if the current image map input is the base case:
// a single image named defaultImage
func isDefaultImg(imgs map[string]string) bool {
	_, ok := imgs[defaultImgName]
	return ok && len(imgs) == 1
}