/*
   Copyright 2023 Docker Compose CLI 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 (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"net/http"
	"path/filepath"
	"slices"
	"time"

	"github.com/containerd/containerd/v2/core/remotes"
	pusherrors "github.com/containerd/containerd/v2/core/remotes/errors"
	"github.com/distribution/reference"
	"github.com/opencontainers/go-digest"
	"github.com/opencontainers/image-spec/specs-go"
	v1 "github.com/opencontainers/image-spec/specs-go/v1"

	"github.com/docker/compose/v5/pkg/api"
)

const (
	// ComposeProjectArtifactType is the OCI 1.1-compliant artifact type value
	// for the generated image manifest.
	ComposeProjectArtifactType = "application/vnd.docker.compose.project"
	// ComposeYAMLMediaType is the media type for each layer (Compose file)
	// in the image manifest.
	ComposeYAMLMediaType = "application/vnd.docker.compose.file+yaml"
	// ComposeEmptyConfigMediaType is a media type used for the config descriptor
	// when doing OCI 1.0-style pushes.
	//
	// The content is always `{}`, the same as a normal empty descriptor, but
	// the specific media type allows clients to fall back to the config media
	// type to recognize the manifest as a Compose project since the artifact
	// type field is not available in OCI 1.0.
	//
	// This is based on guidance from the OCI 1.1 spec:
	//	> Implementers note: artifacts have historically been created without
	// 	> an artifactType field, and tooling to work with artifacts should
	//	> fallback to the config.mediaType value.
	ComposeEmptyConfigMediaType = "application/vnd.docker.compose.config.empty.v1+json"
	// ComposeEnvFileMediaType is the media type for each Env File layer in the image manifest.
	ComposeEnvFileMediaType = "application/vnd.docker.compose.envfile"
)

// clientAuthStatusCodes are client (4xx) errors that are authentication
// related.
var clientAuthStatusCodes = []int{
	http.StatusUnauthorized,
	http.StatusForbidden,
	http.StatusProxyAuthRequired,
}

func DescriptorForComposeFile(path string, content []byte) v1.Descriptor {
	return v1.Descriptor{
		MediaType: ComposeYAMLMediaType,
		Digest:    digest.FromString(string(content)),
		Size:      int64(len(content)),
		Annotations: map[string]string{
			"com.docker.compose.version": api.ComposeVersion,
			"com.docker.compose.file":    filepath.Base(path),
		},
		Data: content,
	}
}

func DescriptorForEnvFile(path string, content []byte) v1.Descriptor {
	return v1.Descriptor{
		MediaType: ComposeEnvFileMediaType,
		Digest:    digest.FromString(string(content)),
		Size:      int64(len(content)),
		Annotations: map[string]string{
			"com.docker.compose.version": api.ComposeVersion,
			"com.docker.compose.envfile": filepath.Base(path),
		},
		Data: content,
	}
}

func PushManifest(ctx context.Context, resolver remotes.Resolver, named reference.Named, layers []v1.Descriptor, ociVersion api.OCIVersion) (v1.Descriptor, error) {
	// Check if we need an extra empty layer for the manifest config
	if ociVersion == api.OCIVersion1_1 || ociVersion == "" {
		err := push(ctx, resolver, named, v1.DescriptorEmptyJSON)
		if err != nil {
			return v1.Descriptor{}, err
		}
	}
	// prepare to push the manifest by pushing the layers
	layerDescriptors := make([]v1.Descriptor, len(layers))
	for i := range layers {
		layerDescriptors[i] = layers[i]
		if err := push(ctx, resolver, named, layers[i]); err != nil {
			return v1.Descriptor{}, err
		}
	}

	if ociVersion != "" {
		// if a version was explicitly specified, use it
		return createAndPushManifest(ctx, resolver, named, layerDescriptors, ociVersion)
	}

	// try to push in the OCI 1.1 format but fallback to OCI 1.0 on 4xx errors
	// (other than auth) since it's most likely the result of the registry not
	// having support
	descriptor, err := createAndPushManifest(ctx, resolver, named, layerDescriptors, api.OCIVersion1_1)
	var pushErr pusherrors.ErrUnexpectedStatus
	if errors.As(err, &pushErr) && isNonAuthClientError(pushErr.StatusCode) {
		// TODO(milas): show a warning here (won't work with logrus)
		return createAndPushManifest(ctx, resolver, named, layerDescriptors, api.OCIVersion1_0)
	}
	return descriptor, err
}

func push(ctx context.Context, resolver remotes.Resolver, ref reference.Named, descriptor v1.Descriptor) error {
	fullRef, err := reference.WithDigest(reference.TagNameOnly(ref), descriptor.Digest)
	if err != nil {
		return err
	}

	return Push(ctx, resolver, fullRef, descriptor)
}

func createAndPushManifest(ctx context.Context, resolver remotes.Resolver, named reference.Named, layers []v1.Descriptor, ociVersion api.OCIVersion) (v1.Descriptor, error) {
	descriptor, toPush, err := generateManifest(layers, ociVersion)
	if err != nil {
		return v1.Descriptor{}, err
	}
	for _, p := range toPush {
		err = push(ctx, resolver, named, p)
		if err != nil {
			return v1.Descriptor{}, err
		}
	}
	return descriptor, nil
}

func isNonAuthClientError(statusCode int) bool {
	if statusCode < 400 || statusCode >= 500 {
		// not a client error
		return false
	}
	return !slices.Contains(clientAuthStatusCodes, statusCode)
}

func generateManifest(layers []v1.Descriptor, ociCompat api.OCIVersion) (v1.Descriptor, []v1.Descriptor, error) {
	var toPush []v1.Descriptor
	var config v1.Descriptor
	var artifactType string
	switch ociCompat {
	case api.OCIVersion1_0:
		// "Content other than OCI container images MAY be packaged using the image manifest.
		// When this is done, the config.mediaType value MUST be set to a value specific to
		// the artifact type or the empty value."
		// Source: https://github.com/opencontainers/image-spec/blob/main/manifest.md#guidelines-for-artifact-usage
		//
		// The `ComposeEmptyConfigMediaType` is used specifically for this purpose:
		// there is no config, and an empty descriptor is used for OCI 1.1 in
		// conjunction with the `ArtifactType`, but for OCI 1.0 compatibility,
		// tooling falls back to the config media type, so this is used to
		// indicate that it's not a container image but custom content.
		configData := []byte("{}")
		config = v1.Descriptor{
			MediaType: ComposeEmptyConfigMediaType,
			Digest:    digest.FromBytes(configData),
			Size:      int64(len(configData)),
			Data:      configData,
		}
		// N.B. OCI 1.0 does NOT support specifying the artifact type, so it's
		//		left as an empty string to omit it from the marshaled JSON
		artifactType = ""
		toPush = append(toPush, config)
	case api.OCIVersion1_1:
		config = v1.DescriptorEmptyJSON
		artifactType = ComposeProjectArtifactType
		toPush = append(toPush, config)
	default:
		return v1.Descriptor{}, nil, fmt.Errorf("unsupported OCI version: %s", ociCompat)
	}

	manifest, err := json.Marshal(v1.Manifest{
		Versioned:    specs.Versioned{SchemaVersion: 2},
		MediaType:    v1.MediaTypeImageManifest,
		ArtifactType: artifactType,
		Config:       config,
		Layers:       layers,
		Annotations: map[string]string{
			"org.opencontainers.image.created": time.Now().Format(time.RFC3339),
		},
	})
	if err != nil {
		return v1.Descriptor{}, nil, err
	}

	manifestDescriptor := v1.Descriptor{
		MediaType: v1.MediaTypeImageManifest,
		Digest:    digest.FromString(string(manifest)),
		Size:      int64(len(manifest)),
		Annotations: map[string]string{
			"com.docker.compose.version": api.ComposeVersion,
		},
		ArtifactType: artifactType,
		Data:         manifest,
	}
	toPush = append(toPush, manifestDescriptor)
	return manifestDescriptor, toPush, nil
}
