package mutate

import (
	"errors"
	"fmt"

	v1 "github.com/google/go-containerregistry/pkg/v1"
	"github.com/google/go-containerregistry/pkg/v1/empty"
	"github.com/google/go-containerregistry/pkg/v1/mutate"

	"edge-infra.dev/pkg/f8n/warehouse"
	"edge-infra.dev/pkg/f8n/warehouse/oci"
	"edge-infra.dev/pkg/f8n/warehouse/oci/layer"
)

// SkipAll is used as a return value from map Fns to indicate that traversal
// should end immediately, returning the input artifact unmodified.
//
//nolint:revive // Sentinel error values for controlling functionality by user can omit Err prefix. See the stdlib fs packages for examples.
var SkipAll = errors.New("end the traversal immediately")

// Map performs a depth-first traversal of the input artifact's graph, visiting
// each node in reverse topological order from the input artifact's root node:
//
// - Dependencies are always visited before the nodes that depend on them:
//   - When the input is a [v1.ImageIndex], [v1.Image] attached directly to the
//     root index with the same name as the index will be visited last, because
//     they represent cluster variants of the package artifact being mapped,
//     and must be visited last. In a standard DFS, those descriptors
//     could be visited on the first pass.
//
// - Functions are called after transformation, in a bottom up fashion.
func Map(a oci.Artifact, fns *Fns) (oci.Artifact, error) {
	switch t := a.(type) {
	case oci.Unwrapper:
		return Map(t.Unwrap(), fns)
	case v1.ImageIndex:
		return fns.mapIdx(t, nil)
	case v1.Image:
		return fns.mapImg(t, nil)
	default:
		return nil, fmt.Errorf("%w: unexpected Artifact type %T", oci.ErrInvalidArtifact, t)
	}
}

// Fns defines functions for each OCI entity to be called during graph traversal
// in [Map] operations. Functions can mutate the input OCI entity by returning
// a new entity. Returning nil causes the entity to be dropped.
//
// The input OCI entity (idx, img, l) is passed in _after_ transforming all of
// it's child nodes. The parent OCI entity passed in is _before_ transformation,
// as it cannot be visited until after all of its child nodes have.
type Fns struct {
	Index func(idx v1.ImageIndex, parent v1.ImageIndex) (v1.ImageIndex, error)
	Image func(img v1.Image, parent v1.ImageIndex) (v1.Image, error)
	Layer func(l layer.Layer, parent v1.Image) (layer.Layer, error)
}

func (f *Fns) mapIdx(idx v1.ImageIndex, parent v1.ImageIndex) (v1.ImageIndex, error) {
	m, err := idx.IndexManifest()
	if err != nil {
		return nil, err
	}

	var (
		rootImgs   []v1.Image
		xArtifacts []oci.Artifact // transformed idx constituents
		changed    = false        // track whether or not children are mutated
	)

	// First pass, visit all dependencies before dependents:
	// - Visit any descriptors referencing images which don't have the same
	//   name as the index
	for _, desc := range m.Manifests {
		switch {
		case desc.MediaType.IsIndex():
			ii, err := idx.ImageIndex(desc.Digest)
			if err != nil {
				return nil, err
			}

			x, err := f.mapIdx(ii, idx)

			switch {
			case errors.Is(err, SkipAll):
				return idx, nil
			case err != nil:
				return nil, err
			case x == nil:
				changed = true
				continue
			}

			changed = changed || (x != ii)
			xArtifacts = append(xArtifacts, x)
		case desc.MediaType.IsImage():
			img, err := idx.Image(desc.Digest)
			if err != nil {
				return nil, err
			}

			// Stash desc for some fun later, this is the descriptor referencing
			// the root nodes image, which needs to be visited last
			if descRefsIdx(m, desc) {
				rootImgs = append(rootImgs, img)
				continue
			}

			// Otherwise, we can go ahead and visit it, since it is a leaf node
			x, err := f.mapImg(img, idx)

			switch {
			case errors.Is(err, SkipAll):
				return idx, nil
			case err != nil:
				return nil, err
			case x == nil:
				changed = true
				continue
			}

			changed = changed || (x != img)
			xArtifacts = append(xArtifacts, x)
		default:
			return nil, fmt.Errorf("%w: unexpected mediaType %s",
				oci.ErrInvalidArtifact, desc.MediaType)
		}
	}

	// Now that we have traversed everything except for the root package image
	// descriptors, we can visit them.
	for _, img := range rootImgs {
		x, err := f.mapImg(img, idx)

		switch {
		case errors.Is(err, SkipAll):
			return idx, nil
		case err != nil:
			return nil, err
		case x == nil:
			changed = true
			continue
		}

		changed = changed || (x != img)
		xArtifacts = append(xArtifacts, x)
	}

	if !changed {
		return f.idxFn(idx, parent)
	}

	// TODO(aw185176): should probably explicitly preserve mediaType here once
	// google/go-containerregistry cuts a release containing fix for empty.Index
	// mediaType
	n, err := AppendManifests(empty.Index, xArtifacts...)
	if err != nil {
		return nil, err
	}

	// Visit the index now that all of its children have been visited
	return f.idxFn(mutate.Annotations(n, m.Annotations).(v1.ImageIndex), parent)
}

func (f *Fns) mapImg(img v1.Image, parent v1.ImageIndex) (v1.Image, error) {
	// short circuit if there is no chance we change the contents of img
	switch {
	case f.Image == nil && f.Layer == nil:
		return img, nil
	case f.Layer == nil:
		return f.imgFn(img, parent)
	}

	var (
		changed = false        // track if image contents were mutated
		xLayers = []v1.Layer{} // transformed layers
	)

	layers, err := layer.FromImage(img)
	if err != nil {
		return nil, err
	}

	for _, layer := range layers {
		x, err := f.layerFn(layer, img)

		switch {
		case errors.Is(err, SkipAll):
			return img, nil
		case err != nil:
			return nil, err
		case x == nil:
			changed = true
			continue
		}

		changed = changed || (x != layer)
		xLayers = append(xLayers, x)
	}

	if !changed {
		return f.imgFn(img, parent)
	}

	// TODO: generalize image mutation, mutate.Image ?
	n, err := ReplaceLayers(img, xLayers...)
	if err != nil {
		return nil, err
	}

	return f.imgFn(n, parent)
}

// idxFn wraps Fns.Index with a nil check, to simplify mapping logic
func (f *Fns) idxFn(idx v1.ImageIndex, parent v1.ImageIndex) (v1.ImageIndex, error) {
	if f.Index != nil {
		return f.Index(idx, parent)
	}
	return idx, nil
}

// imgFn wraps Fns.Image with a nil check, to simplify mapping logic
func (f *Fns) imgFn(img v1.Image, parent v1.ImageIndex) (v1.Image, error) {
	if f.Image != nil {
		return f.Image(img, parent)
	}
	return img, nil
}

// layerFn wraps Fns.Layer with a nil check, to simplify mapping logic
func (f *Fns) layerFn(l layer.Layer, parent v1.Image) (layer.Layer, error) {
	if f.Layer != nil {
		return f.Layer(l, parent)
	}
	return l, nil
}

// descRefsIdx returns true if the input descriptor is referencing the input
// image manifest, e.g., provider variants of a package
func descRefsIdx(idx *v1.IndexManifest, d v1.Descriptor) bool {
	return idx.Annotations[warehouse.AnnotationName] ==
		d.Annotations[warehouse.AnnotationRefName]
}