/*
   Copyright Docker attest authors

   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 oci

import (
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"strings"

	"github.com/google/go-containerregistry/pkg/name"
	v1 "github.com/google/go-containerregistry/pkg/v1"
	"github.com/google/go-containerregistry/pkg/v1/layout"
	"github.com/google/go-containerregistry/pkg/v1/remote"
)

const (
	OCIReferenceTarget            = "org.opencontainers.image.ref.name"
	LocalPrefix                   = "oci://"
	RegistryPrefix                = "docker://"
	OCI                SourceType = "OCI"
	Docker             SourceType = "Docker"
)

type (
	SourceType string
	NamedIndex struct {
		Index v1.ImageIndex
		Name  string
	}
)

type ImageSpecOption func(*ImageSpec) error

type ImageSpec struct {
	// OCI or Docker
	Type SourceType
	// without oci:// or docker:// (name or path)
	Identifier string
	Platform   *v1.Platform
}

func IndexFromPath(path string) (*NamedIndex, error) {
	wrapperIdx, err := layout.ImageIndexFromPath(path)
	if err != nil {
		return nil, fmt.Errorf("failed to load image index: %w", err)
	}

	idxm, err := wrapperIdx.IndexManifest()
	if err != nil {
		return nil, fmt.Errorf("failed to get digest: %w", err)
	}
	imageName := idxm.Manifests[0].Annotations[OCIReferenceTarget]
	idxDigest := idxm.Manifests[0].Digest

	idx, err := wrapperIdx.ImageIndex(idxDigest)
	if err != nil {
		return nil, fmt.Errorf("failed to extract ImageIndex for digest %s: %w", idxDigest.String(), err)
	}
	return &NamedIndex{
		Index: idx,
		Name:  imageName,
	}, nil
}

func IndexFromRemote(ctx context.Context, image string) (*NamedIndex, error) {
	ref, err := name.ParseReference(image)
	if err != nil {
		return nil, fmt.Errorf("failed to parse image reference %s: %w", image, err)
	}

	// Pull the image from the registry
	idx, err := remote.Index(ref, WithOptions(ctx, nil)...)
	if err != nil {
		return nil, fmt.Errorf("failed to pull image %s: %w", image, err)
	}
	return &NamedIndex{
		Index: idx,
		Name:  image,
	}, nil
}

func LoadIndex(ctx context.Context, input *ImageSpec) (*NamedIndex, error) {
	if input.Type == OCI {
		return IndexFromPath(input.Identifier)
	}
	return IndexFromRemote(ctx, input.Identifier)
}

func (i *ImageSpec) ForPlatforms(platform string) ([]*ImageSpec, error) {
	platforms := strings.Split(platform, ",")
	var specs []*ImageSpec
	for _, pStr := range platforms {
		p, err := ParsePlatform(pStr)
		if err != nil {
			return nil, err
		}
		spec := &ImageSpec{
			Type:       i.Type,
			Identifier: i.Identifier,
			Platform:   p,
		}
		specs = append(specs, spec)
	}
	return specs, nil
}

func ParseImageSpec(img string, options ...ImageSpecOption) (*ImageSpec, error) {
	img = strings.TrimSpace(img)
	if strings.Contains(img, ",") {
		return nil, fmt.Errorf("only one image is supported")
	}
	withoutPrefix := strings.TrimPrefix(strings.TrimPrefix(img, LocalPrefix), RegistryPrefix)
	src := &ImageSpec{
		Identifier: withoutPrefix,
	}
	if strings.HasPrefix(img, LocalPrefix) {
		src.Type = OCI
	} else {
		src.Type = Docker
	}
	for _, option := range options {
		err := option(src)
		if err != nil {
			return nil, err
		}
	}
	if src.Platform == nil {
		platform, err := ParsePlatform("")
		if err != nil {
			return nil, err
		}
		src.Platform = platform
	}
	return src, nil
}

func WithPlatform(platform string) ImageSpecOption {
	return func(i *ImageSpec) error {
		if strings.Contains(platform, ",") {
			return fmt.Errorf("only one platform is supported")
		}
		p, err := ParsePlatform(platform)
		if err != nil {
			return err
		}
		i.Platform = p
		return nil
	}
}

func ParseImageSpecs(img string) ([]*ImageSpec, error) {
	outputs := strings.Split(img, ",")
	var sources []*ImageSpec
	for _, output := range outputs {
		src, err := ParseImageSpec(output)
		if err != nil {
			return nil, err
		}
		sources = append(sources, src)
	}
	return sources, nil
}

func WithoutTag(image string) (string, error) {
	if strings.HasPrefix(image, LocalPrefix) {
		return image, nil
	}
	prefix := ""
	if strings.HasPrefix(image, RegistryPrefix) {
		image = strings.TrimPrefix(image, RegistryPrefix)
		prefix = RegistryPrefix
	}
	ref, err := name.ParseReference(image)
	if err != nil {
		return "", err
	}
	repo := ref.Context().Name()
	return prefix + repo, nil
}

type EmptyConfigImage struct {
	v1.Image
}

func (i *EmptyConfigImage) RawConfigFile() ([]byte, error) {
	return []byte("{}"), nil
}

func (i *EmptyConfigImage) Manifest() (*v1.Manifest, error) {
	mf, err := i.Image.Manifest()
	if err != nil {
		return nil, fmt.Errorf("failed to get manifest: %w", err)
	}
	mf.Config = v1.Descriptor{
		MediaType: "application/vnd.oci.empty.v1+json",
		Size:      2,
		Digest:    v1.Hash{Algorithm: "sha256", Hex: "44136fa355b3678a1146ad16f7e8649e94fb4fc21fe77e8310c060f61caaff8a"},
		Data:      []byte("{}"),
	}
	return mf, nil
}

func (i *EmptyConfigImage) RawManifest() ([]byte, error) {
	mf, err := i.Manifest()
	if err != nil {
		return nil, fmt.Errorf("failed to get manifest: %w", err)
	}
	return json.Marshal(mf)
}

func (i *EmptyConfigImage) Digest() (v1.Hash, error) {
	mb, err := i.RawManifest()
	if err != nil {
		return v1.Hash{}, err
	}
	digest, _, err := v1.SHA256(bytes.NewReader(mb))
	return digest, err
}