package image

import (
	"context"
	"fmt"
	"os"
	"sync"
	"time"

	"github.com/containerd/log"
	"github.com/docker/docker/errdefs"
	"github.com/docker/docker/layer"
	"github.com/opencontainers/go-digest"
	"github.com/opencontainers/go-digest/digestset"
	"github.com/pkg/errors"
)

// Store is an interface for creating and accessing images
type Store interface {
	Create(config []byte) (ID, error)
	Get(id ID) (*Image, error)
	Delete(id ID) ([]layer.Metadata, error)
	Search(partialID string) (ID, error)
	SetParent(id ID, parent ID) error
	GetParent(id ID) (ID, error)
	SetLastUpdated(id ID) error
	GetLastUpdated(id ID) (time.Time, error)
	SetBuiltLocally(id ID) error
	IsBuiltLocally(id ID) (bool, error)
	Children(id ID) []ID
	Map() map[ID]*Image
	Heads() map[ID]*Image
	Len() int
}

// LayerGetReleaser is a minimal interface for getting and releasing images.
type LayerGetReleaser interface {
	Get(layer.ChainID) (layer.Layer, error)
	Release(layer.Layer) ([]layer.Metadata, error)
}

type imageMeta struct {
	layer    layer.Layer
	children map[ID]struct{}
}

type store struct {
	sync.RWMutex
	lss       LayerGetReleaser
	images    map[ID]*imageMeta
	fs        StoreBackend
	digestSet *digestset.Set
}

// NewImageStore returns new store object for given set of layer stores
func NewImageStore(fs StoreBackend, lss LayerGetReleaser) (Store, error) {
	is := &store{
		lss:       lss,
		images:    make(map[ID]*imageMeta),
		fs:        fs,
		digestSet: digestset.NewSet(),
	}

	// load all current images and retain layers
	if err := is.restore(); err != nil {
		return nil, err
	}

	return is, nil
}

func (is *store) restore() error {
	// As the code below is run when restoring all images (which can be "many"),
	// constructing the "log.G(ctx).WithFields" is deliberately not "DRY", as the
	// logger is only used for error-cases, and we don't want to do allocations
	// if we don't need it. The "f" type alias is here is just for convenience,
	// and to make the code _slightly_ more DRY. See the discussion on GitHub;
	// https://github.com/moby/moby/pull/44426#discussion_r1059519071
	type f = log.Fields
	err := is.fs.Walk(func(dgst digest.Digest) error {
		img, err := is.Get(ID(dgst))
		if err != nil {
			log.G(context.TODO()).WithFields(f{"digest": dgst, "err": err}).Error("invalid image")
			return nil
		}
		var l layer.Layer
		if chainID := img.RootFS.ChainID(); chainID != "" {
			if err := CheckOS(img.OperatingSystem()); err != nil {
				log.G(context.TODO()).WithFields(f{"chainID": chainID, "os": img.OperatingSystem()}).Error("not restoring image with unsupported operating system")
				return nil
			}
			l, err = is.lss.Get(chainID)
			if err != nil {
				if errors.Is(err, layer.ErrLayerDoesNotExist) {
					log.G(context.TODO()).WithFields(f{"chainID": chainID, "os": img.OperatingSystem(), "err": err}).Error("not restoring image")
					return nil
				}
				return err
			}
		}
		if err := is.digestSet.Add(dgst); err != nil {
			return err
		}

		is.images[ID(dgst)] = &imageMeta{
			layer:    l,
			children: make(map[ID]struct{}),
		}

		return nil
	})
	if err != nil {
		return err
	}

	// Second pass to fill in children maps
	for id := range is.images {
		if parent, err := is.GetParent(id); err == nil {
			if parentMeta := is.images[parent]; parentMeta != nil {
				parentMeta.children[id] = struct{}{}
			}
		}
	}

	return nil
}

func (is *store) Create(config []byte) (ID, error) {
	var img *Image
	img, err := NewFromJSON(config)
	if err != nil {
		return "", err
	}

	// Must reject any config that references diffIDs from the history
	// which aren't among the rootfs layers.
	rootFSLayers := make(map[layer.DiffID]struct{})
	for _, diffID := range img.RootFS.DiffIDs {
		rootFSLayers[diffID] = struct{}{}
	}

	layerCounter := 0
	for _, h := range img.History {
		if !h.EmptyLayer {
			layerCounter++
		}
	}
	if layerCounter > len(img.RootFS.DiffIDs) {
		return "", errdefs.InvalidParameter(errors.New("too many non-empty layers in History section"))
	}

	imageDigest, err := is.fs.Set(config)
	if err != nil {
		return "", errdefs.InvalidParameter(err)
	}

	is.Lock()
	defer is.Unlock()

	imageID := ID(imageDigest)
	if _, exists := is.images[imageID]; exists {
		return imageID, nil
	}

	layerID := img.RootFS.ChainID()

	var l layer.Layer
	if layerID != "" {
		if err := CheckOS(img.OperatingSystem()); err != nil {
			return "", err
		}
		l, err = is.lss.Get(layerID)
		if err != nil {
			return "", errdefs.InvalidParameter(errors.Wrapf(err, "failed to get layer %s", layerID))
		}
	}

	is.images[imageID] = &imageMeta{
		layer:    l,
		children: make(map[ID]struct{}),
	}

	if err = is.digestSet.Add(imageDigest); err != nil {
		delete(is.images, imageID)
		return "", errdefs.InvalidParameter(err)
	}

	return imageID, nil
}

type imageNotFoundError string

func (e imageNotFoundError) Error() string {
	return "No such image: " + string(e)
}

func (imageNotFoundError) NotFound() {}

func (is *store) Search(term string) (ID, error) {
	dgst, err := is.digestSet.Lookup(term)
	if err != nil {
		if errors.Is(err, digestset.ErrDigestNotFound) {
			err = imageNotFoundError(term)
		}
		return "", errors.WithStack(err)
	}
	return ID(dgst), nil
}

func (is *store) Get(id ID) (*Image, error) {
	// todo: Check if image is in images
	// todo: Detect manual insertions and start using them
	config, err := is.fs.Get(id.Digest())
	if err != nil {
		return nil, errdefs.NotFound(err)
	}

	img, err := NewFromJSON(config)
	if err != nil {
		return nil, errdefs.InvalidParameter(err)
	}
	img.computedID = id

	img.Parent, err = is.GetParent(id)
	if err != nil {
		img.Parent = ""
	}

	return img, nil
}

func (is *store) Delete(id ID) ([]layer.Metadata, error) {
	is.Lock()
	defer is.Unlock()

	imgMeta := is.images[id]
	if imgMeta == nil {
		return nil, errdefs.NotFound(fmt.Errorf("unrecognized image ID %s", id.String()))
	}
	_, err := is.Get(id)
	if err != nil {
		return nil, errdefs.NotFound(fmt.Errorf("unrecognized image %s, %v", id.String(), err))
	}
	for cID := range imgMeta.children {
		is.fs.DeleteMetadata(cID.Digest(), "parent")
	}
	if parent, err := is.GetParent(id); err == nil && is.images[parent] != nil {
		delete(is.images[parent].children, id)
	}

	if err := is.digestSet.Remove(id.Digest()); err != nil {
		log.G(context.TODO()).Errorf("error removing %s from digest set: %q", id, err)
	}
	delete(is.images, id)
	is.fs.Delete(id.Digest())

	if imgMeta.layer != nil {
		return is.lss.Release(imgMeta.layer)
	}
	return nil, nil
}

func (is *store) SetParent(id, parentID ID) error {
	is.Lock()
	defer is.Unlock()
	parentMeta := is.images[parentID]
	if parentMeta == nil {
		return errdefs.NotFound(fmt.Errorf("unknown parent image ID %s", parentID.String()))
	}
	if parent, err := is.GetParent(id); err == nil && is.images[parent] != nil {
		delete(is.images[parent].children, id)
	}
	parentMeta.children[id] = struct{}{}
	return is.fs.SetMetadata(id.Digest(), "parent", []byte(parentID))
}

func (is *store) GetParent(id ID) (ID, error) {
	d, err := is.fs.GetMetadata(id.Digest(), "parent")
	if err != nil {
		return "", errdefs.NotFound(err)
	}
	return ID(d), nil // todo: validate?
}

// SetLastUpdated time for the image ID to the current time
func (is *store) SetLastUpdated(id ID) error {
	lastUpdated := []byte(time.Now().Format(time.RFC3339Nano))
	return is.fs.SetMetadata(id.Digest(), "lastUpdated", lastUpdated)
}

// GetLastUpdated time for the image ID
func (is *store) GetLastUpdated(id ID) (time.Time, error) {
	bytes, err := is.fs.GetMetadata(id.Digest(), "lastUpdated")
	if err != nil || len(bytes) == 0 {
		// No lastUpdated time
		return time.Time{}, nil
	}
	return time.Parse(time.RFC3339Nano, string(bytes))
}

// SetBuiltLocally sets whether image can be used as a builder cache
func (is *store) SetBuiltLocally(id ID) error {
	return is.fs.SetMetadata(id.Digest(), "builtLocally", []byte{1})
}

// IsBuiltLocally returns whether image can be used as a builder cache
func (is *store) IsBuiltLocally(id ID) (bool, error) {
	bytes, err := is.fs.GetMetadata(id.Digest(), "builtLocally")
	if err != nil || len(bytes) == 0 {
		if errors.Is(err, os.ErrNotExist) {
			err = nil
		}
		return false, err
	}
	return bytes[0] == 1, nil
}

func (is *store) Children(id ID) []ID {
	is.RLock()
	defer is.RUnlock()

	return is.children(id)
}

func (is *store) children(id ID) []ID {
	var ids []ID
	if is.images[id] != nil {
		for id := range is.images[id].children {
			ids = append(ids, id)
		}
	}
	return ids
}

func (is *store) Heads() map[ID]*Image {
	return is.imagesMap(false)
}

func (is *store) Map() map[ID]*Image {
	return is.imagesMap(true)
}

func (is *store) imagesMap(all bool) map[ID]*Image {
	is.RLock()
	defer is.RUnlock()

	images := make(map[ID]*Image)

	for id := range is.images {
		if !all && len(is.children(id)) > 0 {
			continue
		}
		img, err := is.Get(id)
		if err != nil {
			log.G(context.TODO()).Errorf("invalid image access: %q, error: %q", id, err)
			continue
		}
		images[id] = img
	}
	return images
}

func (is *store) Len() int {
	is.RLock()
	defer is.RUnlock()
	return len(is.images)
}
