package imagetools

import (
	"bytes"
	"context"
	"encoding/json"
	"maps"
	"os"
	"strings"

	"github.com/containerd/containerd/v2/core/content"
	"github.com/containerd/containerd/v2/core/images"
	"github.com/containerd/containerd/v2/core/remotes"
	"github.com/containerd/errdefs"
	"github.com/containerd/platforms"
	"github.com/moby/buildkit/client/ociindex"
	"github.com/moby/buildkit/exporter/containerimage/exptypes"
	"github.com/moby/buildkit/util/attestation"
	"github.com/moby/buildkit/util/contentutil"
	"github.com/opencontainers/go-digest"
	"github.com/opencontainers/image-spec/specs-go"
	ocispecs "github.com/opencontainers/image-spec/specs-go/v1"
	"github.com/pkg/errors"
	"golang.org/x/sync/errgroup"
)

const (
	artifactTypeAttestationManifest = "application/vnd.docker.attestation.manifest.v1+json"
	artifactTypeCosignSignature     = "application/vnd.dev.cosign.artifact.sig.v1+json"
	artifactTypeSigstoreBundle      = "application/vnd.dev.sigstore.bundle.v0.3+json"
)

var supportedArtifactTypes = map[string]struct{}{
	artifactTypeAttestationManifest: {},
	artifactTypeCosignSignature:     {},
	artifactTypeSigstoreBundle:      {},
}

type Source struct {
	Desc ocispecs.Descriptor
	Ref  *Location
}

func (r *Resolver) Combine(ctx context.Context, srcs []*Source, ann map[exptypes.AnnotationKey]string, preferIndex bool, platforms []ocispecs.Platform) ([]byte, ocispecs.Descriptor, []DescWithSource, error) {
	dt, desc, srcMap, err := r.combine(ctx, srcs, ann, preferIndex)
	if err != nil {
		return nil, ocispecs.Descriptor{}, nil, err
	}
	dt, desc, mfstsWithSource, err := r.filterPlatforms(ctx, dt, desc, srcMap, platforms)
	if err != nil {
		return nil, ocispecs.Descriptor{}, nil, err
	}
	return dt, desc, mfstsWithSource, nil
}

func (r *Resolver) combine(ctx context.Context, srcs []*Source, ann map[exptypes.AnnotationKey]string, preferIndex bool) ([]byte, ocispecs.Descriptor, map[digest.Digest]*Source, error) {
	eg, ctx := errgroup.WithContext(ctx)

	dts := make([][]byte, len(srcs))
	for i := range dts {
		func(i int) {
			eg.Go(func() error {
				dt, err := r.GetDescriptor(ctx, srcs[i].Ref, srcs[i].Desc)
				if err != nil {
					return err
				}
				dts[i] = dt

				if srcs[i].Desc.MediaType == "" {
					mt, err := detectMediaType(dt)
					if err != nil {
						return err
					}
					srcs[i].Desc.MediaType = mt
				}

				mt := srcs[i].Desc.MediaType

				switch mt {
				case images.MediaTypeDockerSchema2Manifest, ocispecs.MediaTypeImageManifest:
					p := srcs[i].Desc.Platform
					if srcs[i].Desc.Platform == nil {
						p = &ocispecs.Platform{}
					}
					if p.OS == "" || p.Architecture == "" {
						if err := r.loadPlatform(ctx, srcs[i].Ref, p, dt); err != nil {
							return err
						}
					}
					srcs[i].Desc.Platform = p
				case images.MediaTypeDockerSchema1Manifest:
					return errors.Errorf("schema1 manifests are not allowed in manifest lists")
				}

				return nil
			})
		}(i)
	}

	if err := eg.Wait(); err != nil {
		return nil, ocispecs.Descriptor{}, nil, err
	}

	// on single source, return original bytes
	if len(srcs) == 1 && len(ann) == 0 {
		switch srcs[0].Desc.MediaType {
		// if the source is already an image index or manifest list, there is no need to consider the value
		// of preferIndex since if set to true then the source is already in the preferred format, and if false
		// it doesn't matter since we're not going to split it into separate manifests
		case images.MediaTypeDockerSchema2ManifestList, ocispecs.MediaTypeImageIndex:
			srcMap := map[digest.Digest]*Source{
				srcs[0].Desc.Digest: srcs[0],
			}
			return dts[0], srcs[0].Desc, srcMap, nil
		default:
			if !preferIndex {
				return dts[0], srcs[0].Desc, nil, nil
			}
		}
	}

	indexes := map[digest.Digest]int{}
	sources := map[digest.Digest]*Source{}
	descs := make([]ocispecs.Descriptor, 0, len(srcs))

	addDesc := func(d ocispecs.Descriptor, src *Source) {
		idx, ok := indexes[d.Digest]
		if ok {
			old := descs[idx]
			if old.MediaType == "" {
				old.MediaType = d.MediaType
			}
			if d.Platform != nil {
				old.Platform = d.Platform
			}
			if old.Annotations == nil {
				old.Annotations = map[string]string{}
			}
			maps.Copy(old.Annotations, d.Annotations)
			descs[idx] = old
		} else {
			indexes[d.Digest] = len(descs)
			descs = append(descs, d)
		}
		sources[d.Digest] = src
	}

	for i, src := range srcs {
		switch src.Desc.MediaType {
		case images.MediaTypeDockerSchema2ManifestList, ocispecs.MediaTypeImageIndex:
			var mfst ocispecs.Index
			if err := json.Unmarshal(dts[i], &mfst); err != nil {
				return nil, ocispecs.Descriptor{}, nil, errors.WithStack(err)
			}
			for _, d := range mfst.Manifests {
				addDesc(d, src)
			}
		default:
			addDesc(src.Desc, src)
		}
	}

	dockerMfsts := 0
	for _, desc := range descs {
		if strings.HasPrefix(desc.MediaType, "application/vnd.docker.") {
			dockerMfsts++
		}
	}

	var mt string
	if dockerMfsts == len(descs) {
		// all manifests are Docker types, use Docker manifest list
		mt = images.MediaTypeDockerSchema2ManifestList
	} else {
		// otherwise, use OCI index
		mt = ocispecs.MediaTypeImageIndex
	}

	// annotations are only allowed on OCI indexes
	indexAnnotation := make(map[string]string)
	if mt == ocispecs.MediaTypeImageIndex {
		for k, v := range ann {
			switch k.Type {
			case exptypes.AnnotationIndex:
				indexAnnotation[k.Key] = v
			case exptypes.AnnotationManifestDescriptor:
				for i := range descs {
					if descs[i].Annotations == nil {
						descs[i].Annotations = map[string]string{}
					}
					if k.Platform == nil || k.PlatformString() == platforms.Format(*descs[i].Platform) {
						descs[i].Annotations[k.Key] = v
					}
				}
			case exptypes.AnnotationManifest, "":
				return nil, ocispecs.Descriptor{}, nil, errors.Errorf("%q annotations are not supported yet", k.Type)
			case exptypes.AnnotationIndexDescriptor:
				return nil, ocispecs.Descriptor{}, nil, errors.Errorf("%q annotations are invalid while creating an image", k.Type)
			}
		}
	}

	idxBytes, err := json.MarshalIndent(ocispecs.Index{
		MediaType: mt,
		Versioned: specs.Versioned{
			SchemaVersion: 2,
		},
		Manifests:   descs,
		Annotations: indexAnnotation,
	}, "", "  ")
	if err != nil {
		return nil, ocispecs.Descriptor{}, nil, errors.Wrap(err, "failed to marshal index")
	}

	return idxBytes, ocispecs.Descriptor{
		MediaType: mt,
		Size:      int64(len(idxBytes)),
		Digest:    digest.FromBytes(idxBytes),
	}, sources, nil
}

func (r *Resolver) Push(ctx context.Context, ref *Location, desc ocispecs.Descriptor, dt []byte) error {
	ctx = remotes.WithMediaTypeKeyPrefix(ctx, "application/vnd.in-toto+json", "intoto")
	if ref.IsOCILayout() {
		if err := ref.ValidateTargetDigest(desc.Digest); err != nil {
			return err
		}
		return r.pushOCILayout(ctx, ref, desc, dt)
	}

	fullRef, err := ref.WithDigest(desc.Digest)
	if err != nil {
		return errors.Wrapf(err, "failed to combine ref %s with digest %s", ref, desc.Digest)
	}
	p, err := r.registryResolver().Pusher(ctx, fullRef.String())
	if err != nil {
		return err
	}
	cw, err := p.Push(ctx, desc)
	if err != nil {
		if errdefs.IsAlreadyExists(err) {
			return nil
		}
		return err
	}

	err = content.Copy(ctx, cw, bytes.NewReader(dt), desc.Size, desc.Digest)
	if errdefs.IsAlreadyExists(err) {
		return nil
	}
	return err
}

func (r *Resolver) Copy(ctx context.Context, src *Source, dest *Location) error {
	ctx = remotes.WithMediaTypeKeyPrefix(ctx, "application/vnd.in-toto+json", "intoto")
	ctx = remotes.WithMediaTypeKeyPrefix(ctx, "application/vnd.oci.empty.v1+json", "empty")

	desc := src.Desc
	desc.Annotations = maps.Clone(desc.Annotations)
	if desc.Annotations == nil {
		desc.Annotations = make(map[string]string)
	}
	if src.Ref.IsRegistry() {
		desc.Annotations["containerd.io/distribution.source."+src.Ref.Named().Name()] = src.Ref.Named().Name()
	}

	provider, err := r.providerForLocation(src.Ref)
	if err != nil {
		return err
	}
	ingester, err := r.ingesterForLocation(dest)
	if err != nil {
		return err
	}

	referrers := &referrersProvider{base: referrersFunc(func(ctx context.Context, subject ocispecs.Descriptor) ([]ocispecs.Descriptor, error) {
		descs, err := r.FetchReferrers(ctx, src.Ref, subject.Digest)
		if err != nil {
			return nil, err
		}
		var filtered []ocispecs.Descriptor
		for _, d := range descs {
			if _, ok := supportedArtifactTypes[d.ArtifactType]; ok {
				filtered = append(filtered, d)
			}
		}
		return filtered, nil
	})}

	err = contentutil.CopyChain(ctx, ingester, provider, desc, contentutil.WithReferrers(referrers))
	if err != nil {
		return err
	}
	if dest.IsOCILayout() {
		for subject, descs := range referrers.refs {
			r.ociReferrers.record(dest.OCILayout().Path, subject, descs)
		}
	}
	return nil
}

func (r *Resolver) loadPlatform(ctx context.Context, loc *Location, p2 *ocispecs.Platform, dt []byte) error {
	var manifest ocispecs.Manifest
	if err := json.Unmarshal(dt, &manifest); err != nil {
		return errors.WithStack(err)
	}

	dt, err := r.GetDescriptor(ctx, loc, manifest.Config)
	if err != nil {
		return err
	}

	var p ocispecs.Platform
	if err := json.Unmarshal(dt, &p); err != nil {
		return errors.WithStack(err)
	}

	p = platforms.Normalize(p)

	if p2.Architecture == "" {
		p2.Architecture = p.Architecture
		if p2.Variant == "" {
			p2.Variant = p.Variant
		}
	}
	if p2.OS == "" {
		p2.OS = p.OS
	}

	return nil
}

type referrersFunc func(ctx context.Context, desc ocispecs.Descriptor) ([]ocispecs.Descriptor, error)

func (f referrersFunc) Referrers(ctx context.Context, desc ocispecs.Descriptor) ([]ocispecs.Descriptor, error) {
	return f(ctx, desc)
}

type DescWithSource struct {
	ocispecs.Descriptor
	Source *Source
}

func (r *Resolver) filterPlatforms(ctx context.Context, dt []byte, desc ocispecs.Descriptor, srcMap map[digest.Digest]*Source, plats []ocispecs.Platform) ([]byte, ocispecs.Descriptor, []DescWithSource, error) {
	matcher := platforms.Any(plats...)
	if len(plats) == 0 {
		matcher = platforms.All
	}

	if !images.IsIndexType(desc.MediaType) {
		var mfst ocispecs.Manifest
		if err := json.Unmarshal(dt, &mfst); err != nil {
			return nil, ocispecs.Descriptor{}, nil, errors.Wrapf(err, "failed to parse manifest")
		}
		if desc.Platform == nil {
			return nil, ocispecs.Descriptor{}, nil, errors.Errorf("cannot filter platforms from a manifest without platform information")
		}
		if !matcher.Match(*desc.Platform) {
			return nil, ocispecs.Descriptor{}, nil, errors.Errorf("input platform %s does not match any of the provided platforms", platforms.Format(*desc.Platform))
		}
		return dt, desc, nil, nil
	}

	var idx ocispecs.Index
	if err := json.Unmarshal(dt, &idx); err != nil {
		return nil, ocispecs.Descriptor{}, nil, errors.Wrapf(err, "failed to parse index")
	}

	var manifestMap = map[digest.Digest]ocispecs.Descriptor{}
	for _, m := range idx.Manifests {
		manifestMap[m.Digest] = m
	}
	var references = map[digest.Digest]ocispecs.Descriptor{}
	var matchedManifests = map[digest.Digest]struct{}{}
	for _, m := range idx.Manifests {
		if m.Platform == nil || matcher.Match(*m.Platform) {
			matchedManifests[m.Digest] = struct{}{}
		}
		if refType, ok := m.Annotations[attestation.DockerAnnotationReferenceType]; ok && refType == attestation.DockerAnnotationReferenceTypeDefault {
			dgstStr, ok := m.Annotations[attestation.DockerAnnotationReferenceDigest]
			if !ok {
				continue
			}
			dgst, err := digest.Parse(dgstStr)
			if err != nil {
				continue
			}
			subject, ok := manifestMap[dgst]
			if !ok {
				continue
			}
			if subject.Platform == nil || matcher.Match(*subject.Platform) {
				references[m.Digest] = subject
			}
		}
	}

	var mfsts []ocispecs.Descriptor
	var mfstsWithSource []DescWithSource

	for _, m := range idx.Manifests {
		_, isRef := references[m.Digest]
		if isRef || m.Platform == nil || matcher.Match(*m.Platform) {
			src, ok := srcMap[m.Digest]
			if !ok {
				defaultSource, ok := srcMap[desc.Digest]
				if !ok {
					return nil, ocispecs.Descriptor{}, nil, errors.Errorf("internal error: no source found for %s", m.Digest)
				}
				src = defaultSource
			}
			mfsts = append(mfsts, m)
			mfstsWithSource = append(mfstsWithSource, DescWithSource{
				Descriptor: m,
				Source:     src,
			})
		}
	}

	if len(mfsts) == 0 {
		return nil, ocispecs.Descriptor{}, nil, errors.Errorf("none of the manifests match the provided platforms")
	}

	// try to pull in attestation manifest via referrer if one exists
	addedRef := false
	for d := range matchedManifests {
		if _, ok := references[d]; ok { // manifest itself is already attestation
			continue
		}
		hasRef := false
		for _, subject := range references {
			if subject.Digest == d {
				hasRef = true
				break
			}
		}
		if hasRef {
			continue
		}
		src, ok := srcMap[d]
		if !ok {
			defaultSource, ok := srcMap[desc.Digest]
			if !ok {
				return nil, ocispecs.Descriptor{}, nil, errors.Errorf("internal error: no source found for %s", d)
			}
			src = defaultSource
		}
		refs, err := r.FetchReferrers(ctx, src.Ref, d, remotes.WithReferrerArtifactTypes(artifactTypeAttestationManifest))
		if err != nil {
			if errors.Is(err, errdefs.ErrNotFound) {
				continue
			}
			return nil, ocispecs.Descriptor{}, nil, err
		}
		for _, ref := range refs {
			if _, ok := references[ref.Digest]; ok {
				continue
			}
			ref.Platform = &ocispecs.Platform{
				OS: "unknown", Architecture: "unknown",
			}
			if ref.Annotations == nil {
				ref.Annotations = map[string]string{}
			}
			ref.Annotations[attestation.DockerAnnotationReferenceType] = attestation.DockerAnnotationReferenceTypeDefault
			ref.Annotations[attestation.DockerAnnotationReferenceDigest] = d.String()
			ref.ArtifactType = ""
			mfsts = append(mfsts, ref)
			addedRef = true
			break
		}
	}

	if len(mfsts) == len(idx.Manifests) && !addedRef {
		// all platforms matched, no need to rewrite index
		return dt, desc, mfstsWithSource, nil
	}

	idx.Manifests = mfsts
	idxBytes, err := json.MarshalIndent(&idx, "", "  ")
	if err != nil {
		return nil, ocispecs.Descriptor{}, nil, errors.Wrap(err, "failed to marshal index")
	}

	desc = ocispecs.Descriptor{
		MediaType:   desc.MediaType,
		Size:        int64(len(idxBytes)),
		Digest:      digest.FromBytes(idxBytes),
		Annotations: desc.Annotations,
	}

	return idxBytes, desc, mfstsWithSource, nil
}

type referrersProvider struct {
	base referrersFunc
	refs map[digest.Digest][]ocispecs.Descriptor
}

func (r *referrersProvider) Referrers(ctx context.Context, desc ocispecs.Descriptor) ([]ocispecs.Descriptor, error) {
	out, err := r.base(ctx, desc)
	if err != nil {
		return nil, err
	}
	out = dedupeDescriptors(out)
	if r.refs == nil {
		r.refs = map[digest.Digest][]ocispecs.Descriptor{}
	}
	r.refs[desc.Digest] = dedupeDescriptors(append(r.refs[desc.Digest], out...))
	return out, nil
}

func dedupeDescriptors(descs []ocispecs.Descriptor) []ocispecs.Descriptor {
	if len(descs) < 2 {
		return descs
	}
	seen := make(map[digest.Digest]struct{}, len(descs))
	out := descs[:0]
	for _, desc := range descs {
		if _, ok := seen[desc.Digest]; ok {
			continue
		}
		seen[desc.Digest] = struct{}{}
		out = append(out, desc)
	}
	return out
}

func (r *Resolver) ingesterForLocation(loc *Location) (content.Ingester, error) {
	if loc.IsRegistry() {
		p, err := r.registryResolver().Pusher(context.TODO(), loc.Name())
		if err != nil {
			return nil, err
		}
		return contentutil.FromPusher(p), nil
	}
	return r.localStore(loc.OCILayout().Path)
}

func (r *Resolver) providerForLocation(loc *Location) (content.Provider, error) {
	if loc.IsRegistry() {
		f, err := r.registryResolver().Fetcher(context.TODO(), loc.String())
		if err != nil {
			return nil, err
		}
		return contentutil.FromFetcher(f), nil
	}
	return r.localStore(loc.OCILayout().Path)
}

func (r *Resolver) pushOCILayout(ctx context.Context, ref *Location, desc ocispecs.Descriptor, dt []byte) error {
	if err := os.MkdirAll(ref.OCILayout().Path, 0o755); err != nil {
		return err
	}
	store, err := r.localStore(ref.OCILayout().Path)
	if err != nil {
		return err
	}
	w, err := store.Writer(ctx, content.WithRef(desc.Digest.String()), content.WithDescriptor(desc))
	if err != nil {
		if !errdefs.IsAlreadyExists(err) {
			return err
		}
	} else {
		err = content.Copy(ctx, w, bytes.NewReader(dt), desc.Size, desc.Digest)
		if err != nil && !errdefs.IsAlreadyExists(err) {
			return err
		}
	}

	idx := ociindex.NewStoreIndex(ref.OCILayout().Path)
	switch {
	case ref.Digest() != "":
		if err := idx.Put(desc); err != nil {
			return err
		}
	case ref.Tag() != "":
		if err := idx.Put(desc, ociindex.Tag(ref.Tag())); err != nil {
			return err
		}
	default:
		if err := idx.Put(desc, ociindex.Tag("latest")); err != nil {
			return err
		}
	}
	return writePendingOCILayoutReferrers(ctx, r.ociReferrers.take(ref.OCILayout().Path), r.GetDescriptor, idx, ref)
}

func detectMediaType(dt []byte) (string, error) {
	var mfst struct {
		MediaType string          `json:"mediaType"`
		Config    json.RawMessage `json:"config"`
		FSLayers  []string        `json:"fsLayers"`
	}

	if err := json.Unmarshal(dt, &mfst); err != nil {
		return "", errors.WithStack(err)
	}

	if mfst.MediaType != "" {
		return mfst.MediaType, nil
	}
	if mfst.Config != nil {
		return images.MediaTypeDockerSchema2Manifest, nil
	}
	if len(mfst.FSLayers) > 0 {
		return images.MediaTypeDockerSchema1Manifest, nil
	}

	return images.MediaTypeDockerSchema2ManifestList, nil
}