package layer

import (
	"fmt"
	"sort"

	"github.com/fluxcd/pkg/ssa"
	v1 "github.com/google/go-containerregistry/pkg/v1"
	"github.com/google/go-containerregistry/pkg/v1/partial"
	"github.com/google/go-containerregistry/pkg/v1/static"

	wh "edge-infra.dev/pkg/f8n/warehouse"
	"edge-infra.dev/pkg/k8s/unstructured"
)

// ErrInvalidLayer can occur if the v1.Layer is not a valid Warehouse layer,
// per the Warehouse spec.
var ErrInvalidLayer = fmt.Errorf("OCI layer is not a valid Warehouse layer")

// Layer is a v1.Layer with additional accessors and Warehouse-specific
// functionality. It wraps a v1.Layer in a way that it can preserve annotations
// when partial.Descriptor() is called with it. This is required because we use
// static.Layer, which does not implement Descriptor()
type Layer interface {
	v1.Layer

	// Descriptor implements partial.withDescriptor so that partial.Descriptor(layer)
	// returns the expected annotations.
	Descriptor() (*v1.Descriptor, error)

	// Key returns a stable identifier for this layer based on its type.
	//
	// If a layer is a runtime capability, the runtime capability name is returned.
	// Otherwise, the layer's key is it's type
	Key() string

	// Type returns the Warehouse layer type.
	Type() Type

	// Unstructured returns the layer's contents as untyped K8s objects that can be
	// applied to the K8s API server
	Unstructured() ([]*unstructured.Unstructured, error)

	// Annotations returns the annotations present for this layer.
	Annotations() map[string]string
}

// layer provides a basic concrete implementation of Layer. It can be used as
// is or as the basis for another implementation (e.g., mutate.Layer)
type layer struct {
	v1.Layer
	t     Type
	annos map[string]string
}

// layer implements Layer
var _ Layer = &layer{}

// New creates a Warehouse layer for the given type and contents. Contents are
// uncompressed.
func New(t Type, contents []byte, opts ...Option) (Layer, error) {
	if err := t.IsValid(); err != nil {
		return nil, err
	}
	o := makeOptions(opts...)
	if o.cap != "" && t != Runtime {
		return nil, fmt.Errorf(
			"layer.New: ForCapability cannot be used with non-Runtime layer types",
		)
	}

	l := &layer{
		static.NewLayer(contents, YAML),
		t,
		map[string]string{wh.AnnotationLayerType: string(t)},
	}

	for k, v := range o.annos {
		l.annos[k] = v
	}

	return l, nil
}

// Descriptor implements partial.withDescriptor so that partial.Descriptor(layer)
// returns the expected annotations.
func (l *layer) Descriptor() (*v1.Descriptor, error) {
	d, err := l.Digest()
	if err != nil {
		return nil, err
	}

	s, err := l.Size()
	if err != nil {
		return nil, err
	}

	mt, err := l.MediaType()
	if err != nil {
		return nil, err
	}

	if l.annos == nil {
		l.annos = make(map[string]string)
	}

	return &v1.Descriptor{
		Digest:      d,
		Size:        s,
		MediaType:   mt,
		Annotations: l.annos,
	}, nil
}

func (l *layer) Key() string {
	if c, ok := l.annos[wh.AnnotationLayerRuntimeCapability]; ok {
		return c
	}
	return string(l.t)
}

func (l *layer) Type() Type {
	return l.t
}

func (l *layer) Unstructured() ([]*unstructured.Unstructured, error) {
	r, err := l.Uncompressed()
	if err != nil {
		return nil, err
	}
	return ssa.ReadObjects(r)
}

func (l *layer) Annotations() map[string]string {
	return l.annos
}

// Layer sorting
// TODO(aw185176): test layer sorting
type layers []Layer

var _ sort.Interface = layers{}

func (l layers) Len() int      { return len(l) }
func (l layers) Swap(i, j int) { l[i], l[j] = l[j], l[i] }
func (l layers) Less(i, j int) bool {
	switch {
	// everything is less than base runtime
	case l[j].Key() == Runtime.String():
		return false
	case l[i].Key() == Runtime.String():
		return true
	// everything other than base runtime layer is less than infra
	case l[i].Key() == Infra.String():
		return true
	default:
		return false
	}
}

// Sort sorts the input layers by their key in place.
func Sort(ll []Layer) {
	sort.Stable(layers(ll))
}

// FromImage parses Warehouse layers from a standard OCI image. Any layers that
// aren't valid Warehouse layers are dropped.
func FromImage(img v1.Image) ([]Layer, error) {
	imgLayers, err := img.Layers()
	if err != nil {
		return nil, err
	}

	ll := make([]Layer, 0, len(imgLayers))
	for _, il := range imgLayers {
		l, err := FromOCILayer(il)
		if err == nil {
			ll = append(ll, l)
		}
	}
	return ll, nil
}

// FromOCILayer creates a Warehouse Layer from a standard v1.Layer. If it isn't
// a valid OCI v1.Layer, ErrInvalidLayer is returned.
func FromOCILayer(l v1.Layer) (Layer, error) {
	d, err := partial.Descriptor(l)
	if err != nil {
		return nil, err
	}

	if d.Annotations[wh.AnnotationLayerType] == "" {
		return nil, fmt.Errorf("%w: missing %s annotation",
			ErrInvalidLayer, wh.AnnotationLayerType)
	}

	return &layer{
		Layer: l,
		t:     Type(d.Annotations[wh.AnnotationLayerType]),
		annos: d.Annotations,
	}, nil
}