// Copyright 2019 Google LLC All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//    http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package mutate

import (
	"bytes"
	"encoding/json"
	"errors"
	"sync"

	v1 "github.com/google/go-containerregistry/pkg/v1"
	"github.com/google/go-containerregistry/pkg/v1/partial"
	"github.com/google/go-containerregistry/pkg/v1/stream"
	"github.com/google/go-containerregistry/pkg/v1/types"
)

type image struct {
	base v1.Image
	adds []Addendum

	computed        bool
	configFile      *v1.ConfigFile
	manifest        *v1.Manifest
	annotations     map[string]string
	mediaType       *types.MediaType
	configMediaType *types.MediaType
	diffIDMap       map[v1.Hash]v1.Layer
	digestMap       map[v1.Hash]v1.Layer
	subject         *v1.Descriptor

	sync.Mutex
}

var _ v1.Image = (*image)(nil)

func (i *image) MediaType() (types.MediaType, error) {
	if i.mediaType != nil {
		return *i.mediaType, nil
	}
	return i.base.MediaType()
}

// isImageConfig reports whether the media type is a Docker or OCI image config.
func isImageConfig(mt types.MediaType) bool {
	switch mt {
	case types.DockerConfigJSON, types.OCIConfigJSON:
		return true
	default:
		return false
	}
}

func (i *image) compute() error {
	i.Lock()
	defer i.Unlock()

	// Don't re-compute if already computed.
	if i.computed {
		return nil
	}

	m, err := i.base.Manifest()
	if err != nil {
		return err
	}
	manifest := m.DeepCopy()

	diffIDMap := make(map[v1.Hash]v1.Layer)
	digestMap := make(map[v1.Hash]v1.Layer)

	// Determine effective config media type (user override takes precedence).
	cfgMediaType := manifest.Config.MediaType
	if i.configMediaType != nil {
		cfgMediaType = *i.configMediaType
	}

	imageConfig := isImageConfig(cfgMediaType)

	var configFile *v1.ConfigFile
	if i.configFile != nil {
		configFile = i.configFile
	} else {
		cf, err := i.base.ConfigFile()
		if err != nil {
			return err
		}
		configFile = cf.DeepCopy()
	}

	// For image configs, update RootFS.DiffIDs and History from added layers.
	// For artifacts, skip this: the config has no rootfs or history fields.
	if imageConfig {
		diffIDs := configFile.RootFS.DiffIDs
		history := configFile.History

		for _, add := range i.adds {
			history = append(history, add.History)
			if add.Layer != nil {
				diffID, err := add.Layer.DiffID()
				if err != nil {
					return err
				}
				diffIDs = append(diffIDs, diffID)
				diffIDMap[diffID] = add.Layer
			}
		}

		configFile.RootFS.DiffIDs = diffIDs
		configFile.History = history
	}

	manifestLayers := manifest.Layers
	for _, add := range i.adds {
		if add.Layer == nil {
			continue
		}

		desc, err := partial.Descriptor(add.Layer)
		if err != nil {
			return err
		}

		if len(add.Annotations) != 0 {
			desc.Annotations = add.Annotations
		}
		if len(add.URLs) != 0 {
			desc.URLs = add.URLs
		}
		if add.MediaType != "" {
			desc.MediaType = add.MediaType
		}

		manifestLayers = append(manifestLayers, *desc)
		digestMap[desc.Digest] = add.Layer
	}
	manifest.Layers = manifestLayers

	// For image configs, re-marshal the config and update the manifest digest.
	// For artifacts, preserve the original config blob as-is to avoid
	// corrupting the digest via re-marshaling.
	if imageConfig {
		rcfg, err := json.Marshal(configFile)
		if err != nil {
			return err
		}
		d, sz, err := v1.SHA256(bytes.NewBuffer(rcfg))
		if err != nil {
			return err
		}
		manifest.Config.Digest = d
		manifest.Config.Size = sz

		if m.Config.Data != nil {
			manifest.Config.Data = rcfg
		}
	}

	if i.configMediaType != nil {
		manifest.Config.MediaType = *i.configMediaType
	}
	if i.mediaType != nil {
		manifest.MediaType = *i.mediaType
	}
	if i.annotations != nil {
		if manifest.Annotations == nil {
			manifest.Annotations = map[string]string{}
		}
		for k, v := range i.annotations {
			manifest.Annotations[k] = v
		}
	}
	manifest.Subject = i.subject

	i.configFile = configFile
	i.manifest = manifest
	i.diffIDMap = diffIDMap
	i.digestMap = digestMap
	i.computed = true
	return nil
}

// Layers returns the ordered collection of filesystem layers that comprise this image.
// The order of the list is oldest/base layer first, and most-recent/top layer last.
func (i *image) Layers() ([]v1.Layer, error) {
	if err := i.compute(); errors.Is(err, stream.ErrNotComputed) || (i.manifest != nil && !isImageConfig(i.manifest.Config.MediaType)) {
		// Stream not yet consumed, or non-image OCI artifact (RootFS.DiffIDs
		// is empty so partial.DiffIDs returns nothing). Fall back to the base
		// layers plus any added layers.
		layers, err := i.base.Layers()
		if err != nil {
			return nil, err
		}
		for _, add := range i.adds {
			if add.Layer != nil {
				layers = append(layers, add.Layer)
			}
		}
		return layers, nil
	} else if err != nil {
		return nil, err
	}

	// Walk manifest layer descriptors by digest rather than rootfs diff
	// IDs. Two layers can legitimately share a diff ID — same uncompressed
	// content, different compression — and produce distinct digests. The
	// manifest preserves the per-occurrence digest; LayerByDiffID does not,
	// which previously caused duplicate-diff-ID layers to collapse to a
	// single entry in the returned slice and break blob upload for
	// downstream pushers (see #2034).
	ls := make([]v1.Layer, 0, len(i.manifest.Layers))
	for _, desc := range i.manifest.Layers {
		l, err := i.LayerByDigest(desc.Digest)
		if err != nil {
			return nil, err
		}
		ls = append(ls, l)
	}
	return ls, nil
}

// ConfigName returns the hash of the image's config file.
func (i *image) ConfigName() (v1.Hash, error) {
	if err := i.compute(); err != nil {
		return v1.Hash{}, err
	}
	return partial.ConfigName(i)
}

// ConfigFile returns this image's config file.
func (i *image) ConfigFile() (*v1.ConfigFile, error) {
	if err := i.compute(); err != nil {
		return nil, err
	}
	return i.configFile.DeepCopy(), nil
}

// RawConfigFile returns the serialized bytes of ConfigFile().
// For non-image OCI artifacts, returns the original raw config to preserve
// the config blob digest.
func (i *image) RawConfigFile() ([]byte, error) {
	if err := i.compute(); err != nil {
		return nil, err
	}
	// If the manifest config is not a standard image config, return the
	// original raw bytes to avoid corrupting the digest via re-marshaling.
	if i.manifest != nil && !isImageConfig(i.manifest.Config.MediaType) {
		return i.base.RawConfigFile()
	}
	return json.Marshal(i.configFile)
}

// Digest returns the sha256 of this image's manifest.
func (i *image) Digest() (v1.Hash, error) {
	if err := i.compute(); err != nil {
		return v1.Hash{}, err
	}
	return partial.Digest(i)
}

// Size implements v1.Image.
func (i *image) Size() (int64, error) {
	if err := i.compute(); err != nil {
		return -1, err
	}
	return partial.Size(i)
}

// Manifest returns this image's Manifest object.
func (i *image) Manifest() (*v1.Manifest, error) {
	if err := i.compute(); err != nil {
		return nil, err
	}
	return i.manifest.DeepCopy(), nil
}

// RawManifest returns the serialized bytes of Manifest()
func (i *image) RawManifest() ([]byte, error) {
	if err := i.compute(); err != nil {
		return nil, err
	}
	return json.Marshal(i.manifest)
}

// LayerByDigest returns a Layer for interacting with a particular layer of
// the image, looking it up by "digest" (the compressed hash).
func (i *image) LayerByDigest(h v1.Hash) (v1.Layer, error) {
	if cn, err := i.ConfigName(); err != nil {
		return nil, err
	} else if h == cn {
		return partial.ConfigLayer(i)
	}
	if layer, ok := i.digestMap[h]; ok {
		return layer, nil
	}
	return i.base.LayerByDigest(h)
}

// LayerByDiffID is an analog to LayerByDigest, looking up by "diff id"
// (the uncompressed hash).
func (i *image) LayerByDiffID(h v1.Hash) (v1.Layer, error) {
	if layer, ok := i.diffIDMap[h]; ok {
		return layer, nil
	}
	return i.base.LayerByDiffID(h)
}

func validate(adds []Addendum) error {
	for _, add := range adds {
		if add.Layer == nil && !add.History.EmptyLayer {
			return errors.New("unable to add a nil layer to the image")
		}
	}
	return nil
}
