// Package oci contains OCI functionality and types that are used to support
// Warehouse specific implementations in other packages. It is used to enforce
// Warehouse semantics on top of OCI semantics in a reusable way.
//
// Behavior that is specific to types of Warehouse packages (e.g., Pallets),
// should be placed in that package, not here.
package oci

import (
	"bytes"
	"fmt"

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

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

// Artifact represents the intersection between v1.Image and v1.ImageIndex by
// embedding the two partial interfaces required for common OCI operations in
// google/go-containerregistry. This interface allows us to pass around warehouse
// artifacts without concerning callers about whether or not the artifact is an
// v1.Image or v1.ImageIndex.
//
// Artifacts can be used to represent both Pallets and Cluster package types,
// as well as any future Warehouse package types, as long as they embed this
// interface.
type Artifact interface {
	// Describable represents anything that we can produce a descriptor for,
	// e.g., what gets embedded in in v1.Image and v1.ImageIndex to reference
	// other OCI entities.
	// Specifically, that means things that implement partial.Describable can be
	// appended or composed with existing artifacts via the mutate package.
	partial.Describable

	// WithRawManifest is the lowest common denominator which can be used to
	// implement v1.Image and v1.ImageIndex. Nearly all implementations for those
	// interfaces in google/go-containerregistry embed this interface, allowing
	// us to easily use those packages when working with artifacts.
	partial.WithRawManifest
}

// Unwrapper allows access to the underlying OCI artifact via the Unwrap method.
type Unwrapper interface {
	// Unwrap returns the underlying oci.Artifact, allowing for matching based on
	// ggcr interface types (v1.ImageIndex, v1.Image, etc)
	Unwrap() Artifact
}

// Annotate adds a series of annotation maps to an OCI artifact and returns
// the annotated artifact.
// Note that OCI artifacts are immutable: this function is really creating a new
// artifact that is the input + the annotations.
//
// Does not do key collision checks.  Must be casted back to original type (add example)
func Annotate(a partial.WithRawManifest, annos ...map[string]string) partial.WithRawManifest {
	m := map[string]string{}
	for _, anno := range annos {
		for k, v := range anno {
			m[k] = v
		}
	}

	return mutate.Annotations(a, m)
}

// Annotations returns the OCI manifest annotations for the input artifact after
// inferring the concrete type.
func Annotations(a Artifact) (map[string]string, error) {
	switch a := a.(type) {
	case Unwrapper:
		return Annotations(a.Unwrap())
	case v1.Image:
		// lazily evaluate manifest to avoid doing it twice if input is Unwrapper
		raw, err := a.RawManifest()
		if err != nil {
			return nil, fmt.Errorf("failed to load oci manifest from artifact: %w", err)
		}
		manifest, err := v1.ParseManifest(bytes.NewReader(raw))
		if err != nil {
			return nil, fmt.Errorf("failed to parse image manifest: %w", err)
		}
		return manifest.Annotations, nil
	case v1.ImageIndex:
		// lazily evaluate manifest to avoid doing it twice if input is Unwrapper
		raw, err := a.RawManifest()
		if err != nil {
			return nil, fmt.Errorf("failed to load oci manifest from artifact: %w", err)
		}
		manifest, err := v1.ParseIndexManifest(bytes.NewReader(raw))
		if err != nil {
			return nil, fmt.Errorf("failed to parse image index manifest: %w", err)
		}
		return manifest.Annotations, nil
	default:
		return nil, fmt.Errorf("%w", ErrInvalidArtifact)
	}
}

// Contains returns true if the v1.ImageIndex embeds a descriptor referencing the
// Artifact, after inferring the concrete type. The descriptor's ref name
// annotation must match the Artifacts name annotation and the digests must match.
func Contains(idx v1.ImageIndex, a Artifact) (bool, error) {
	digest, err := a.Digest()
	if err != nil {
		return false, err
	}
	annos, err := Annotations(a)
	if err != nil {
		return false, err
	}

	idxm, err := idx.IndexManifest()
	if err != nil {
		return false, err
	}

	for _, m := range idxm.Manifests {
		if m.Annotations[warehouse.AnnotationRefName] ==
			annos[warehouse.AnnotationName] && digest == m.Digest {
			return true, nil
		}
	}

	return false, nil
}

// IsPackage determines if the v1.Image is a standalone package or not by checking
// if it is embedded into the Artifact and has the same name as the
// Artifact. e.g., v1.Images named shoot will be embedded into a v1.ImageIndex
// named shoot in order to represent different provider variants, or if shoot has
// dependencies.
//
// If the Artifact is nil, the Image is assumed to be a package. Anything other
// than nil or an ImageIndex will return an error, as it means you have goofed
// up pretty thoroughly.
//
// If an ImageIndex contains an Image, that represents either an ImageIndex for
// a package with provider-specific variants (shoot:gke, shoot:sds) or a package
// with dependencies. This is a useful signal while walking artifacts from root
// nodes. If you only care about walking the root nodes of each individual package
// in the graph, you can use this function to skip walking Images
func IsPackage(img v1.Image, parent Artifact) (bool, error) {
	imgAnnos, err := Annotations(img)
	if err != nil {
		return false, err
	}

	switch a := parent.(type) {
	case Unwrapper:
		return IsPackage(img, a.Unwrap())
	case v1.ImageIndex:
		manifest, err := a.IndexManifest()
		if err != nil {
			return false, err
		}

		return imgAnnos[warehouse.AnnotationName] !=
			manifest.Annotations[warehouse.AnnotationName], nil

		// confirm that idx embeds image explicitly
		// this level of safety isn't explicitly required in the context of
		// using oci.Walk, since oci.Walk is traversing the descriptors it is
		// guaranteed that the parent embeds the input img. however, that
		// isn't the case if some other client was using IsPackage() outside of
		// the conext of walk. if that use case ever shows up, uncomment this
		// code:
		// refs, err := Contains(a, img)
		// if err != nil {
		// 	return false, err
		// }
		// return !refs, nil
	case nil:
		return true, nil
	default:
		return false, fmt.Errorf("non-ImageIndex type %T was provided as parent", a)
	}
}

// IsComposite returns true if the Artifact is a v1.ImageIndex which only contains
// embedded descriptors referencing other v1.ImageIndexes or standalone packages
// that are v1.Images. This means that the index is a composite index that only
// exists to compose multiple distinct package together.
//
// If the Artifact is anything but an ImageIndex, it exits gracefully with false.
func IsComposite(a Artifact) (bool, error) {
	switch a := a.(type) {
	case Unwrapper:
		return IsComposite(a.Unwrap())
	case v1.ImageIndex:
		manifest, err := a.IndexManifest()
		if err != nil {
			return false, err
		}
		for _, m := range manifest.Manifests {
			if !m.MediaType.IsIndex() &&
				m.Annotations[warehouse.AnnotationRefName] ==
					manifest.Annotations[warehouse.AnnotationName] {
				return false, nil
			}
		}
		return true, nil
	default:
		return false, nil
	}
}

// TODO: keep? doc?
func ArtifactFromIdx(idx v1.ImageIndex, d v1.Descriptor) (Artifact, error) {
	switch {
	case d.MediaType.IsImage():
		return idx.Image(d.Digest)
	case d.MediaType.IsIndex():
		return idx.ImageIndex(d.Digest)
	default:
		return nil, fmt.Errorf("%w: unexpected mediaType %s",
			ErrInvalidArtifact, d.MediaType)
	}
}