package containerimage

import (
	"context"
	"encoding/base64"
	"encoding/json"
	"fmt"
	"maps"
	"sort"
	"strconv"
	"strings"

	"github.com/containerd/containerd/v2/core/content"
	"github.com/containerd/containerd/v2/core/images"
	"github.com/containerd/containerd/v2/core/leases"
	"github.com/containerd/containerd/v2/core/remotes/docker"
	remoteserrors "github.com/containerd/containerd/v2/core/remotes/errors"
	"github.com/containerd/containerd/v2/pkg/epoch"
	"github.com/containerd/containerd/v2/pkg/labels"
	"github.com/containerd/containerd/v2/pkg/rootfs"
	cerrdefs "github.com/containerd/errdefs"
	"github.com/containerd/platforms"
	"github.com/moby/buildkit/cache"
	cacheconfig "github.com/moby/buildkit/cache/config"
	"github.com/moby/buildkit/client"
	"github.com/moby/buildkit/exporter"
	"github.com/moby/buildkit/exporter/containerimage/exptypes"
	"github.com/moby/buildkit/session"
	"github.com/moby/buildkit/snapshot"
	"github.com/moby/buildkit/solver/llbsolver/compat"
	"github.com/moby/buildkit/util/compression"
	"github.com/moby/buildkit/util/contentutil"
	"github.com/moby/buildkit/util/errutil"
	"github.com/moby/buildkit/util/leaseutil"
	"github.com/moby/buildkit/util/progress"
	"github.com/moby/buildkit/util/push"
	digest "github.com/opencontainers/go-digest"
	"github.com/opencontainers/image-spec/identity"
	ocispecs "github.com/opencontainers/image-spec/specs-go/v1"
	"github.com/pkg/errors"
	"golang.org/x/sync/errgroup"
)

const (
	// keyUnsafeInternalStoreAllowIncomplete should only be used for tests. This option allows exporting image to the image store
	// as well as lacking some blobs in the content store. Some integration tests for lazyref behaviour depends on this option.
	// Ignored when store=false.
	keyUnsafeInternalStoreAllowIncomplete = "unsafe-internal-store-allow-incomplete"
)

type Opt struct {
	SessionManager *session.Manager
	ImageWriter    *ImageWriter
	Images         images.Store
	RegistryHosts  docker.RegistryHosts
	LeaseManager   leases.Manager
}

type imageExporter struct {
	opt Opt
}

// New returns a new containerimage exporter instance that supports exporting
// to an image store and pushing the image to registry.
// This exporter supports following values in returned kv map:
// - containerimage.digest - The digest of the root manifest for the image.
func New(opt Opt) (exporter.Exporter, error) {
	im := &imageExporter{opt: opt}
	return im, nil
}

func (e *imageExporter) Resolve(ctx context.Context, id int, opt map[string]string) (exporter.ExporterInstance, error) {
	i := &imageExporterInstance{
		imageExporter: e,
		id:            id,
		attrs:         opt,
		opts: ImageCommitOpts{
			RefCfg: cacheconfig.RefConfig{
				Compression: compression.New(compression.Default),
			},
			ForceInlineAttestations: true,
		},
		store: true,
	}

	opt, err := i.opts.Load(ctx, opt)
	if err != nil {
		return nil, err
	}

	for k, v := range opt {
		switch exptypes.ImageExporterOptKey(k) {
		case exptypes.OptKeyPush:
			if v == "" {
				i.push = true
				continue
			}
			b, err := strconv.ParseBool(v)
			if err != nil {
				return nil, errors.Wrapf(err, "non-bool value specified for %s", k)
			}
			i.push = b
		case exptypes.OptKeyPushByDigest:
			if v == "" {
				i.pushByDigest = true
				continue
			}
			b, err := strconv.ParseBool(v)
			if err != nil {
				return nil, errors.Wrapf(err, "non-bool value specified for %s", k)
			}
			i.pushByDigest = b
		case exptypes.OptKeyInsecure:
			if v == "" {
				i.insecure = true
				continue
			}
			b, err := strconv.ParseBool(v)
			if err != nil {
				return nil, errors.Wrapf(err, "non-bool value specified for %s", k)
			}
			i.insecure = b
		case exptypes.OptKeyUnpack:
			if v == "" {
				i.unpack = true
				continue
			}
			b, err := strconv.ParseBool(v)
			if err != nil {
				return nil, errors.Wrapf(err, "non-bool value specified for %s", k)
			}
			i.unpack = b
		case exptypes.OptKeyStore:
			if v == "" {
				i.store = true
				continue
			}
			b, err := strconv.ParseBool(v)
			if err != nil {
				return nil, errors.Wrapf(err, "non-bool value specified for %s", k)
			}
			i.store = b
		case keyUnsafeInternalStoreAllowIncomplete:
			if v == "" {
				i.storeAllowIncomplete = true
				continue
			}
			b, err := strconv.ParseBool(v)
			if err != nil {
				return nil, errors.Wrapf(err, "non-bool value specified for %s", k)
			}
			i.storeAllowIncomplete = b
		case exptypes.OptKeyDanglingPrefix:
			i.danglingPrefix = v
		case exptypes.OptKeyDanglingEmptyOnly:
			if v == "" {
				i.danglingEmptyOnly = true
				continue
			}
			b, err := strconv.ParseBool(v)
			if err != nil {
				return nil, errors.Wrapf(err, "non-bool value specified for %s", k)
			}
			i.danglingEmptyOnly = b
		case exptypes.OptKeyNameCanonical:
			if v == "" {
				i.nameCanonical = true
				continue
			}
			b, err := strconv.ParseBool(v)
			if err != nil {
				return nil, errors.Wrapf(err, "non-bool value specified for %s", k)
			}
			i.nameCanonical = b
		default:
			if i.meta == nil {
				i.meta = make(map[string][]byte)
			}
			i.meta[k] = []byte(v)
		}
	}
	return i, nil
}

type imageExporterInstance struct {
	*imageExporter
	id    int
	attrs map[string]string

	opts                 ImageCommitOpts
	push                 bool
	pushByDigest         bool
	unpack               bool
	store                bool
	storeAllowIncomplete bool
	insecure             bool
	nameCanonical        bool
	danglingPrefix       string
	danglingEmptyOnly    bool
	meta                 map[string][]byte
}

func (e *imageExporterInstance) ID() int {
	return e.id
}

func (e *imageExporterInstance) Name() string {
	return "exporting to image"
}

func (e *imageExporterInstance) Config() *exporter.Config {
	return exporter.NewConfigWithCompression(e.opts.RefCfg.Compression)
}

func (e *imageExporterInstance) Type() string {
	return client.ExporterImage
}

func (e *imageExporterInstance) Attrs() map[string]string {
	return e.attrs
}

func (e *imageExporterInstance) Export(ctx context.Context, src *exporter.Source, buildInfo exporter.ExportBuildInfo) (_ map[string]string, _ exporter.FinalizeFunc, descref exporter.DescriptorReference, err error) {
	src = src.Clone()
	if src.Metadata == nil {
		src.Metadata = make(map[string][]byte)
	}
	maps.Copy(src.Metadata, e.meta)

	opts := e.opts
	as, _, err := ParseAnnotations(src.Metadata)
	if err != nil {
		return nil, nil, nil, err
	}
	opts.Annotations = opts.Annotations.Merge(as)
	opts.SetOCITypesDefault(DefaultOCITypes(buildInfo.CompatibilityVersion, src))
	if err := opts.Validate(); err != nil {
		return nil, nil, nil, err
	}

	ctx, done, err := leaseutil.WithLease(ctx, e.opt.LeaseManager, leaseutil.MakeTemporary)
	if err != nil {
		return nil, nil, nil, err
	}
	// On success, we create descref which holds the lease's done function.
	// The solver will release descref after recording the descriptor in build
	// history. On error (descref is nil), we release the lease here.
	defer func() {
		if descref == nil {
			done(context.WithoutCancel(ctx))
		}
	}()

	desc, err := e.opt.ImageWriter.Commit(ctx, src, buildInfo.SessionID, buildInfo.InlineCache, &opts, buildInfo.CompatibilityVersion, e.Type())
	if err != nil {
		return nil, nil, nil, err
	}

	resp := make(map[string]string)

	if n, ok := src.Metadata["image.name"]; e.opts.ImageName == "*" && ok {
		e.opts.ImageName = string(n)
	}

	nameCanonical := e.nameCanonical
	if e.danglingPrefix != "" && (!e.danglingEmptyOnly || e.opts.ImageName == "") {
		danglingImageName := e.danglingPrefix + "@" + desc.Digest.String()
		if e.opts.ImageName != "" {
			e.opts.ImageName += "," + danglingImageName
		} else {
			e.opts.ImageName = danglingImageName
			nameCanonical = false
		}
	}

	// Collect names for finalize callback to push
	var namesToPush []string

	if e.opts.ImageName != "" {
		targetNames := strings.SplitSeq(e.opts.ImageName, ",")
		for targetName := range targetNames {
			if e.opt.Images != nil && e.store {
				tagDone := progress.OneOff(ctx, "naming to "+targetName)

				// imageClientCtx is used for propagating the epoch to e.opt.Images.Update() and e.opt.Images.Create().
				//
				// Ideally, we should be able to propagate the epoch via images.Image.CreatedAt.
				// However, due to a bug of containerd, we are temporarily stuck with this workaround.
				// https://github.com/containerd/containerd/issues/8322
				imageClientCtx := ctx
				if e.opts.Epoch != nil && e.opts.Epoch.Value != nil {
					imageClientCtx = epoch.WithSourceDateEpoch(imageClientCtx, e.opts.Epoch.Value)
				}
				img := images.Image{
					Target: *desc,
					// CreatedAt in images.Images is ignored due to a bug of containerd.
					// See the comment lines for imageClientCtx.
				}

				sfx := []string{""}
				if nameCanonical && !strings.ContainsRune(targetName, '@') {
					sfx = append(sfx, "@"+desc.Digest.String())
				}
				for _, sfx := range sfx {
					img.Name = targetName + sfx
					for { // handle possible race between Update and Create
						if _, err := e.opt.Images.Update(imageClientCtx, img); err != nil {
							if !errors.Is(err, cerrdefs.ErrNotFound) {
								return nil, nil, nil, tagDone(err)
							}

							if _, err := e.opt.Images.Create(imageClientCtx, img); err != nil {
								if !errors.Is(err, cerrdefs.ErrAlreadyExists) {
									return nil, nil, nil, tagDone(err)
								}
								continue
							}
						}
						break
					}
				}
				tagDone(nil)

				if e.unpack {
					if opts.RewriteTimestamp {
						// e.unpackImage cannot be used because src ref does not point to the rewritten image
						// /
						// TODO: change e.unpackImage so that it takes Result[Remote] as parameter.
						// https://github.com/moby/buildkit/pull/4057#discussion_r1324106088
						return nil, nil, nil, errors.New("exporter option \"rewrite-timestamp\" conflicts with \"unpack\"")
					}
					if err := e.unpackImage(ctx, img, src, session.NewGroup(buildInfo.SessionID)); err != nil {
						return nil, nil, nil, err
					}
				}

				if !e.storeAllowIncomplete {
					var refs []cache.ImmutableRef
					if src.Ref != nil {
						refs = append(refs, src.Ref)
					}
					for _, ref := range src.Refs {
						if ref == nil {
							continue
						}
						refs = append(refs, ref)
					}
					eg, ctx := errgroup.WithContext(ctx)
					for _, ref := range refs {
						eg.Go(func() error {
							remotes, err := ref.GetRemotes(ctx, false, e.opts.RefCfg, false, session.NewGroup(buildInfo.SessionID))
							if err != nil {
								return err
							}
							remote := remotes[0]
							if unlazier, ok := remote.Provider.(cache.Unlazier); ok {
								if err := unlazier.Unlazy(ctx); err != nil {
									return err
								}
							}
							return nil
						})
					}
					if err := eg.Wait(); err != nil {
						return nil, nil, nil, err
					}
				}
			}
			// Collect names for pushing in finalize
			if e.push {
				namesToPush = append(namesToPush, targetName)
			}
		}
		resp[exptypes.ExporterImageNameKey] = e.opts.ImageName
	}

	resp[exptypes.ExporterImageDigestKey] = desc.Digest.String()
	if v, ok := desc.Annotations[exptypes.ExporterConfigDigestKey]; ok {
		resp[exptypes.ExporterImageConfigDigestKey] = v
		delete(desc.Annotations, exptypes.ExporterConfigDigestKey)
	}

	dtdesc, err := json.Marshal(desc)
	if err != nil {
		return nil, nil, nil, err
	}
	resp[exptypes.ExporterImageDescriptorKey] = base64.StdEncoding.EncodeToString(dtdesc)

	// Create descref so descriptor is recorded in build history.
	// Transfer lease ownership to descref - caller releases after finalize.
	descref = NewDescriptorReference(*desc, done)

	if len(namesToPush) == 0 {
		return resp, nil, descref, nil
	}

	// Create finalize callback for pushing
	finalize := func(ctx context.Context) error {
		for _, targetName := range namesToPush {
			err := e.pushImage(ctx, src, buildInfo.SessionID, targetName, desc.Digest)
			if err != nil {
				var statusErr remoteserrors.ErrUnexpectedStatus
				if errors.As(err, &statusErr) {
					err = errutil.WithDetails(err)
				}
				return errors.Wrapf(err, "failed to push %v", targetName)
			}
		}
		return nil
	}

	return resp, finalize, descref, nil
}

func (e *imageExporterInstance) pushImage(ctx context.Context, src *exporter.Source, sessionID string, targetName string, dgst digest.Digest) error {
	var refs []cache.ImmutableRef
	if src.Ref != nil {
		refs = append(refs, src.Ref)
	}
	for _, ref := range src.Refs {
		if ref == nil {
			continue
		}
		refs = append(refs, ref)
	}

	annotations := map[digest.Digest]map[string]string{}
	mprovider := contentutil.NewMultiProvider(e.opt.ImageWriter.ContentStore())
	for _, ref := range refs {
		remotes, err := ref.GetRemotes(ctx, false, e.opts.RefCfg, false, session.NewGroup(sessionID))
		if err != nil {
			return err
		}
		remote := remotes[0]
		for _, desc := range remote.Descriptors {
			mprovider.Add(desc.Digest, remote.Provider)
			addAnnotations(annotations, desc)
		}
	}
	return push.Push(ctx, e.opt.SessionManager, sessionID, mprovider, e.opt.ImageWriter.ContentStore(), dgst, targetName, e.insecure, e.opt.RegistryHosts, e.pushByDigest, annotations)
}

func (e *imageExporterInstance) unpackImage(ctx context.Context, img images.Image, src *exporter.Source, s session.Group) (err0 error) {
	matcher := platforms.Only(platforms.Normalize(platforms.DefaultSpec()))

	ps, err := exptypes.ParsePlatforms(src.Metadata)
	if err != nil {
		return err
	}
	matching := []exptypes.Platform{}
	for _, p2 := range ps.Platforms {
		if matcher.Match(p2.Platform) {
			matching = append(matching, p2)
		}
	}
	if len(matching) == 0 {
		// current platform was not found, so skip unpacking
		return nil
	}
	sort.SliceStable(matching, func(i, j int) bool {
		return matcher.Less(matching[i].Platform, matching[j].Platform)
	})

	ref, _ := src.FindRef(matching[0].ID)
	if ref == nil {
		// ref has no layers, so nothing to unpack
		return nil
	}

	unpackDone := progress.OneOff(ctx, "unpacking to "+img.Name)
	defer func() {
		unpackDone(err0)
	}()

	var (
		contentStore = e.opt.ImageWriter.ContentStore()
		applier      = e.opt.ImageWriter.Applier()
		snapshotter  = e.opt.ImageWriter.Snapshotter()
	)

	// fetch manifest by default platform
	manifest, err := images.Manifest(ctx, contentStore, img.Target, platforms.Default())
	if err != nil {
		return err
	}

	remotes, err := ref.GetRemotes(ctx, true, e.opts.RefCfg, false, s)
	if err != nil {
		return err
	}
	remote := remotes[0]

	// ensure the content for each layer exists locally in case any are lazy
	if unlazier, ok := remote.Provider.(cache.Unlazier); ok {
		if err := unlazier.Unlazy(ctx); err != nil {
			return err
		}
	}

	layers, err := getLayers(remote.Descriptors, manifest)
	if err != nil {
		return err
	}

	// get containerd snapshotter
	ctrdSnapshotter, release := snapshot.NewContainerdSnapshotter(snapshotter)
	defer release()

	// Compute top chainID so we can add it to the lease before calling ApplyLayers
	// as ApplyLayers may directly return after successful Stat call without applying
	// layer to the lease and causing error if it gets deleted.
	chainID := layersChainID(layers)
	if leaseID, ok := leases.FromContext(ctx); ok {
		r := leases.Resource{
			ID:   chainID.String(),
			Type: "snapshots/" + snapshotter.Name(),
		}
		if err := e.opt.LeaseManager.AddResource(ctx, leases.Lease{ID: leaseID}, r); err != nil {
			return errors.Wrapf(err, "failed to lease snapshot %s", chainID)
		}
	}

	// note that calling ApplyLayer in a loop here as alternative is not safe because
	// single ApplyLayer does not have a safe way to ensure parents are not removed during unpack.
	appliedChainID, err := rootfs.ApplyLayers(ctx, layers, ctrdSnapshotter, applier)
	if err != nil {
		return err
	}
	if appliedChainID != chainID {
		return errors.Errorf("unexpected chain ID mismatch: %s != %s", appliedChainID, chainID)
	}

	var (
		keyGCLabel   = fmt.Sprintf("containerd.io/gc.ref.snapshot.%s", snapshotter.Name())
		valueGCLabel = chainID.String()
	)

	cinfo := content.Info{
		Digest: manifest.Config.Digest,
		Labels: map[string]string{keyGCLabel: valueGCLabel},
	}
	_, err = contentStore.Update(ctx, cinfo, fmt.Sprintf("labels.%s", keyGCLabel))
	return err
}

func getLayers(descs []ocispecs.Descriptor, manifest ocispecs.Manifest) ([]rootfs.Layer, error) {
	if len(descs) != len(manifest.Layers) {
		return nil, errors.Errorf("mismatched image rootfs and manifest layers")
	}

	layers := make([]rootfs.Layer, len(descs))
	for i, desc := range descs {
		layers[i].Diff = ocispecs.Descriptor{
			MediaType: ocispecs.MediaTypeImageLayer,
			Digest:    digest.Digest(desc.Annotations[labels.LabelUncompressed]),
		}
		layers[i].Blob = manifest.Layers[i]
	}
	return layers, nil
}

func layersChainID(layers []rootfs.Layer) digest.Digest {
	chain := make([]digest.Digest, len(layers))
	for i, l := range layers {
		chain[i] = l.Diff.Digest
	}
	return identity.ChainID(chain)
}

func addAnnotations(m map[digest.Digest]map[string]string, desc ocispecs.Descriptor) {
	if desc.Annotations == nil {
		return
	}
	a, ok := m[desc.Digest]
	if !ok {
		m[desc.Digest] = desc.Annotations
		return
	}
	if a == nil {
		a = make(map[string]string)
	}
	maps.Copy(a, desc.Annotations)
}

// DefaultOCITypes returns the default media type behavior for image exports.
func DefaultOCITypes(compatibilityVersion int, src *exporter.Source) bool {
	if compatibilityVersion >= compat.CompatibilityVersion031 {
		return true
	}
	return len(src.Attestations) > 0
}

func NewDescriptorReference(desc ocispecs.Descriptor, release func(context.Context) error) exporter.DescriptorReference {
	return &descriptorReference{
		desc:    desc,
		release: release,
	}
}

type descriptorReference struct {
	desc    ocispecs.Descriptor
	release func(context.Context) error
}

func (d *descriptorReference) Descriptor() ocispecs.Descriptor {
	return d.desc
}

func (d *descriptorReference) Release() error {
	return d.release(context.TODO())
}
