package cmp

import (
	v1 "github.com/google/go-containerregistry/pkg/v1"

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

// Diff returns the list of hashes for blobs that are in a but not in any of bb.
func Diff(a oci.Artifact, bb ...oci.Artifact) ([]v1.Hash, error) {
	// Initial diff is entirety of a. As we encounter hashes traversing bb, we
	// remove matching elements from the diff.
	diff, err := toMap(a)
	if err != nil {
		return nil, err
	}
	// Calculate digest of a, if we encounter the same digest then we know a has
	// no unique blobs and can exit early.
	ad, err := a.Digest()
	if err != nil {
		return nil, err
	}
	hitA := false
	visited := map[v1.Hash]bool{}

	for _, b := range bb {
		// Short circuit entire walk if we have already visited this artifact
		h, err := b.Digest()
		if err != nil {
			return nil, err
		}
		if h == ad {
			return nil, nil
		}
		if visited[h] {
			continue
		}

		if err := walk.Walk(b, &walk.Fns{
			Index: func(ii v1.ImageIndex, _ v1.ImageIndex) error {
				h, err := ii.Digest()
				if err != nil {
					return err
				}

				if h == ad {
					hitA = true
					return walk.SkipAll
				}
				visited[h] = true
				if diff[h] {
					delete(diff, h)
				}
				return nil
			},
			Image: func(img v1.Image, _ v1.ImageIndex) error {
				h, err := img.Digest()
				if err != nil {
					return err
				}
				ch, err := img.ConfigName()
				if err != nil {
					return err
				}

				if h == ad {
					hitA = true
					return walk.SkipAll
				}
				visited[h] = true
				visited[ch] = true

				if diff[h] {
					delete(diff, h)
				}
				if diff[ch] {
					delete(diff, ch)
				}
				return nil
			},
			Layer: func(l layer.Layer, _ v1.Image) error {
				h, err := l.Digest()
				if err != nil {
					return err
				}

				visited[h] = true
				if diff[h] {
					delete(diff, h)
				}
				return nil
			},
		}); err != nil {
			return nil, err
		}

		// If we have encountered every hash in a, then there are no hashes in a
		// that aren't in b.
		if len(diff) == 0 || hitA {
			return nil, nil
		}
	}

	// The remaining elements in diff are the unique hashes in a.
	return mapToSlice(diff), nil
}