package mutate

import (
	"fmt"

	v1 "github.com/google/go-containerregistry/pkg/v1"
	"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/match"
)

// Name mutates the name annotations of Artifact a. If a is an ImageIndex, the
// reference annotations in the embedded descriptors are also mutated.
func Name(a oci.Artifact, newName string) (oci.Artifact, error) {
	switch a := a.(type) {
	case oci.Unwrapper:
		return Name(a.Unwrap(), newName)
	case v1.Image:
		annos, err := oci.Annotations(a)
		if err != nil {
			return nil, err
		}
		annos[warehouse.AnnotationName] = newName
		return oci.Annotate(a, annos).(v1.Image), nil
	case v1.ImageIndex:
		m, err := a.IndexManifest()
		if err != nil {
			return nil, err
		}
		name := m.Annotations[warehouse.AnnotationName]
		if name == "" {
			return nil, fmt.Errorf("mutate.Name: %w: missing name annotation",
				oci.ErrInvalidArtifact)
		}

		// Update name annotation on root artifact before handling children
		m.Annotations[warehouse.AnnotationName] = newName
		a = oci.Annotate(a, m.Annotations).(v1.ImageIndex)

		// Create matcher for finding provider variants
		matcher := match.Annotation(warehouse.AnnotationRefName, name)

		descs, err := match.FindManifests(a, matcher)
		if err != nil {
			return nil, err
		}

		updatedDescs := make([]mutate.IndexAddendum, len(descs))
		for i, d := range descs {
			child, err := oci.ArtifactFromIdx(a, d)
			if err != nil {
				return nil, err
			}
			annos, err := oci.Annotations(child)
			if err != nil {
				return nil, err
			}
			annos[warehouse.AnnotationName] = newName
			child = oci.Annotate(child, annos).(v1.Image)

			d.Annotations[warehouse.AnnotationRefName] = newName

			updatedDescs[i] = mutate.IndexAddendum{
				Add:        child,
				Descriptor: d,
			}
		}

		return mutate.AppendManifests(
			mutate.RemoveManifests(a, matcher),
			updatedDescs...,
		), nil

	default:
		return nil, fmt.Errorf("mutate.Name: %w", oci.ErrInvalidArtifact)
	}
}