// Package layout defines utilities for working with a Pallet OCI Image Layout.
//
// All artifacts written to the layout must have a name annotation representing
// the tag for that artifact, e.g., "shoot:latest".
package layout

import (
	"encoding/json"
	"fmt"
	"io/fs"
	"os"
	"path/filepath"
	"sort"

	"github.com/google/go-containerregistry/pkg/name"
	v1 "github.com/google/go-containerregistry/pkg/v1"
	"github.com/google/go-containerregistry/pkg/v1/empty"
	"github.com/google/go-containerregistry/pkg/v1/layout"
	ociv1 "github.com/opencontainers/image-spec/specs-go/v1"

	wh "edge-infra.dev/pkg/f8n/warehouse"
	"edge-infra.dev/pkg/f8n/warehouse/oci"
	"edge-infra.dev/pkg/f8n/warehouse/oci/cmp"
	"edge-infra.dev/pkg/f8n/warehouse/oci/match"
	whname "edge-infra.dev/pkg/f8n/warehouse/oci/name"
)

// TODO: implement functions for tagging images in the layout and retrieving
// images by tag in the layout? current impl loses name information (since
// name is combined with tag to make a valid ref for the descriptor annotation)
// and doesn't allow for multiple tags of the same artifact. it would also
// be nice to know when something was written to the warehouse

// Path represents an OCI Image Layout on disk. It is used for storing and accessing
// packages locally.
type Path struct {
	layout.Path
}

// New creates a new Path at the specified filepath. If an Image Layout doesn't
// already exist at that path, a new one is created.
func New(path string) (*Path, error) {
	p, err := resolvePath(path)
	if err != nil {
		return nil, err
	}
	return &Path{p}, nil
}

func Clear(path string) (*Path, error) {
	p, err := layout.Write(path, empty.Index)
	if err != nil {
		return nil, fmt.Errorf("failed to create new index at %s: %w", path, err)
	}
	return &Path{p}, nil
}

// Get retrieves an OCI artifact from the layout for the ref, based on what kind
// of ref it is. If the ref is a name.Tag, it is looked up using the OCI reference
// annotation defined in the OCI Image spec. If it is a name.Digest, the digest
// is searched for directly.
func (p *Path) Get(ref name.Reference) (oci.Artifact, error) {
	idx, err := p.ImageIndex()
	if err != nil {
		return nil, fmt.Errorf("failed to load layout index: %w", err)
	}

	if ref, ok := ref.(name.Digest); ok {
		hash, err := v1.NewHash(ref.DigestStr())
		if err != nil {
			return nil, err
		}
		desc, err := match.FindManifest(idx, match.Digests(hash))
		if err != nil {
			return nil, fmt.Errorf("failed to locate %s in warehouse cache: %w",
				ref, err)
		}
		return oci.ArtifactFromIdx(idx, desc)
	}

	desc, err := match.FindManifest(idx, RefMatcher(ref))
	if err != nil {
		return nil, fmt.Errorf("failed to locate %s in warehouse cache: %w",
			ref, err)
	}
	return oci.ArtifactFromIdx(idx, desc)
}

// Append adds the artifact to the image layout with correct warehouse annotations
// for fetching based on the name and tag.
func (p *Path) Append(ref name.Reference, a oci.Artifact) error {
	if a, ok := a.(oci.Unwrapper); ok {
		// Short-circuit recursing with unwrapped artifact if type implements Unwrapper
		return p.Append(ref, a.Unwrap())
	}

	pkgName, err := whname.FromArtifact(a)
	if err != nil {
		return fmt.Errorf("failed to read package name from artifact: %w", err)
	}

	annos := map[string]string{wh.AnnotationRefName: pkgName}
	if t, ok := ref.(name.Tag); ok {
		annos[ociv1.AnnotationRefName] = t.String()
	}

	matcher := RefMatcher(ref)
	if ref, ok := ref.(name.Digest); ok {
		hash, err := v1.NewHash(ref.DigestStr())
		if err != nil {
			return err
		}
		matcher = match.Digests(hash)
	}

	switch a := a.(type) {
	case v1.Image:
		return p.appendImage(a, matcher, layout.WithAnnotations(annos))
	case v1.ImageIndex:
		return p.appendIndex(a, matcher, layout.WithAnnotations(annos))
	default:
		return fmt.Errorf("layout.Append: %w", oci.ErrInvalidArtifact)
	}
}

// Sort improves the reproducibility of the layout's IndexManifest (index.json)
// by sorting the descriptors by digest. The initial index is read in as a GGCR
// v1.IndexManifest. Its descriptors are then sorted by digest. The modified
// index is written back to the same location.
func (p *Path) Sort() error {
	layoutIndex, err := p.ImageIndex()
	if err != nil {
		return fmt.Errorf("failed to retrieve cache index: %w", err)
	}

	idx, err := layoutIndex.IndexManifest()
	if err != nil {
		return fmt.Errorf("failed to unmarshal cache index: %w", err)
	}

	// Sort descriptors by digest
	sort.Slice(idx.Manifests, func(i, j int) bool { return idx.Manifests[i].Digest.String() < idx.Manifests[j].Digest.String() })

	rawIndex, err := json.MarshalIndent(idx, "", "   ")
	if err != nil {
		return fmt.Errorf("failed to marshal sorted cache index: %w", err)
	}

	return p.WriteFile("index.json", rawIndex, os.ModePerm)
}

func (p *Path) appendImage(i v1.Image, m match.Matcher, opts ...layout.Option) error {
	return p.ReplaceImage(i, m, opts...)
}

func (p *Path) appendIndex(ii v1.ImageIndex, m match.Matcher, opts ...layout.Option) error {
	return p.ReplaceIndex(ii, m, opts...)
}

// resolvePath resolves the provided path to a layout.Path which can be
// written to. If an existing Image Layout does not exist at the path, an
// empty one is created and returned.
func resolvePath(path string) (layout.Path, error) {
	p, err := layout.FromPath(path)
	if err != nil {
		p, err = layout.Write(path, empty.Index)
		if err != nil {
			return "", fmt.Errorf("failed to create new index at %s: %w", path, err)
		}
	}
	return p, nil
}

// RefMatcher returns a matcher for looking up the OCI v1 annotation used to
// look up tagged artifacts
func RefMatcher(ref name.Reference) match.Matcher {
	return match.Annotation(ociv1.AnnotationRefName, ref.String())
}

// FromFS reads an Image Layout from the root of fsys. If a layout exists at
// another location in the file system, [fs.Sub] can be used to create an
// appropriate FS:
//
//	f, err := fs.Sub(embeddedFSVar, "layout/root")
//	if err != nil { return err }
//	p, err := layout.FromFS(f)
func FromFS(fsys fs.FS, dst string) (*Path, error) {
	if err := fs.WalkDir(fsys, ".", func(path string, d fs.DirEntry, err error) error {
		if err != nil {
			return err
		}
		if d.IsDir() {
			return nil
		}
		err = os.MkdirAll(filepath.Join(dst, filepath.Dir(path)), 0777)
		if err != nil {
			return err
		}

		data, err := fs.ReadFile(fsys, path)
		if err != nil {
			return fmt.Errorf("failed to read %s from FS: %w", path, err)
		}

		if err := os.WriteFile(filepath.Join(dst, path), data, 0644); err != nil {
			return fmt.Errorf("failed to write %s: %w", filepath.Join(dst, path), err)
		}

		return nil
	}); err != nil {
		return nil, fmt.Errorf("failed to write layout to %s: %w", dst, err)
	}

	p, err := New(dst)
	if err != nil {
		return nil, fmt.Errorf("failed to create Image Layout from %s: %w", dst, err)
	}
	return p, nil
}

// Walk path back to fetch the cached blobs from the given directory
// If error is found in path we default to the home warehouse cache
func (p *Path) cachedBlobs() ([]string, error) {
	dir, err := fs.ReadDir(os.DirFS(filepath.Join(string(p.Path), "blobs/sha256")), ".")
	if err != nil {
		return nil, err
	}
	hashes := make([]string, 0, len(dir))
	for _, d := range dir {
		hashes = append(hashes, d.Name())
	}
	return hashes, nil
}

func (p *Path) Prune() ([]string, error) {
	var pruned []string
	idx, err := p.ImageIndex()
	if err != nil {
		return nil, err
	}
	cached, err := p.cachedBlobs()
	if err != nil {
		return nil, err
	}
	hashes, err := cmp.Hashes(idx)
	if err != nil {
		return nil, err
	}
	blobs := make(map[string]bool, len(hashes))
	for _, i := range hashes {
		blobs[i.Hex] = true
	}
	for _, b := range cached {
		// All blobs that are in the cache that aren't in the image index are pruned
		if !blobs[b] {
			pruned = append(pruned, b)
			b = "sha256:" + b
			hash, err := v1.NewHash(b)
			if err != nil {
				continue
			}
			err = p.RemoveBlob(hash)
			if err != nil {
				return nil, err
			}
		}
	}
	return pruned, nil
}