package store

import (
	"encoding/json"
	"fmt"
	"os"
	"path/filepath"
	"runtime"
	"sort"
	"strings"
	"time"
)

// Load imports the images of a `docker save` archive or an OCI image
// layout into the store (docker load semantics): the outer archive is
// unpacked to a staging directory, each image's layers are flattened
// into a rootfs, and the OCI config becomes the image metadata. The
// archive names its images (RepoTags, or the OCI ref.name annotation);
// reference, when given, overrides that for a single-image archive.
// Returns the references loaded, in short form.
func (s Store) Load(archive, reference string) ([]string, error) {
	archiveAbs, err := filepath.Abs(archive)
	if err != nil {
		return nil, err
	}
	if _, err := os.Stat(archiveAbs); err != nil {
		return nil, fmt.Errorf("archive: %w", err)
	}
	staging, err := os.MkdirTemp(s.imagesRoot(), ".load-")
	if err != nil {
		if mkErr := os.MkdirAll(s.imagesRoot(), 0o755); mkErr != nil {
			return nil, mkErr
		}
		if staging, err = os.MkdirTemp(s.imagesRoot(), ".load-"); err != nil {
			return nil, err
		}
	}
	defer os.RemoveAll(staging)
	unpacked := filepath.Join(staging, "archive")
	if err := untar(archiveAbs, unpacked); err != nil {
		return nil, err
	}
	images, err := readArchiveImages(unpacked)
	if err != nil {
		return nil, fmt.Errorf("%s: %w", archiveAbs, err)
	}
	if reference != "" {
		if len(images) != 1 {
			return nil, fmt.Errorf("%s holds %d images; a name can only be given for one", archiveAbs, len(images))
		}
		images[0].names = []string{reference}
	}
	var loaded []string
	for i, img := range images {
		if len(img.names) == 0 {
			return nil, fmt.Errorf("%s: image %d carries no name (RepoTags); load it with --name", archiveAbs, i+1)
		}
		for _, name := range img.names {
			if _, err := ParseReference(name); err != nil {
				return nil, err
			}
			if _, err := s.Dir(name); err == nil {
				return nil, fmt.Errorf("image %q already exists (werkdock rmi %s to replace it)", name, name)
			}
		}
		dir := filepath.Join(staging, fmt.Sprintf("image-%d", i))
		if err := flattenImage(unpacked, img, dir, archiveAbs); err != nil {
			return nil, err
		}
		// The first name gets the tree, further names of the same image
		// get a hardlinked copy — a tag is a copy in this store (RFC 0004).
		sources := []string{dir}
		for j := 1; j < len(img.names); j++ {
			copy := filepath.Join(staging, fmt.Sprintf("image-%d-%d", i, j))
			if err := copyTree(dir, copy); err != nil {
				return nil, err
			}
			sources = append(sources, copy)
		}
		for j, name := range img.names {
			ref, _ := ParseReference(name)
			if err := writeMeta(sources[j], ImageMeta{Name: ref.Short(), Source: archiveAbs}.merge(img.config)); err != nil {
				return nil, err
			}
			if err := s.place(sources[j], ref); err != nil {
				return nil, err
			}
			loaded = append(loaded, ref.Short())
		}
	}
	return loaded, nil
}

// archiveImage is one image of an archive: where its config and layers
// are (paths relative to the unpacked archive) and what it is called.
type archiveImage struct {
	configPath string
	layers     []string
	names      []string
	config     ociConfig
}

// ociConfig is the OCI image configuration blob, as far as werkdock
// reads it.
type ociConfig struct {
	Architecture string      `json:"architecture"`
	OS           string      `json:"os"`
	Created      time.Time   `json:"created"`
	Config       ImageConfig `json:"config"`
}

func (m ImageMeta) merge(c ociConfig) ImageMeta {
	m.Architecture, m.OS, m.Config = c.Architecture, c.OS, c.Config
	m.CreatedAt = c.Created
	if m.CreatedAt.IsZero() {
		m.CreatedAt = time.Now().UTC()
	}
	if m.Architecture == "" {
		m.Architecture = runtime.GOARCH
	}
	if m.OS == "" {
		m.OS = runtime.GOOS
	}
	return m
}

// readArchiveImages understands the docker save manifest (manifest.json,
// written by every docker version, also beside the OCI layout newer
// ones produce) and, failing that, a plain OCI image layout.
func readArchiveImages(dir string) ([]archiveImage, error) {
	if raw, err := os.ReadFile(filepath.Join(dir, "manifest.json")); err == nil {
		return readDockerManifest(dir, raw)
	}
	if _, err := os.Stat(filepath.Join(dir, "index.json")); err == nil {
		return readOCILayout(dir)
	}
	return nil, fmt.Errorf("neither a docker save archive (manifest.json) nor an OCI image layout (index.json); a plain rootfs archive is imported with: werkdock import")
}

func readDockerManifest(dir string, raw []byte) ([]archiveImage, error) {
	var manifest []struct {
		Config   string   `json:"Config"`
		RepoTags []string `json:"RepoTags"`
		Layers   []string `json:"Layers"`
	}
	if err := json.Unmarshal(raw, &manifest); err != nil {
		return nil, fmt.Errorf("manifest.json: %w", err)
	}
	var images []archiveImage
	for _, m := range manifest {
		img := archiveImage{configPath: m.Config, layers: m.Layers, names: m.RepoTags}
		if err := img.readConfig(dir); err != nil {
			return nil, err
		}
		images = append(images, img)
	}
	return images, nil
}

// OCI descriptors and manifests, the fields werkdock needs.
type ociDescriptor struct {
	MediaType   string            `json:"mediaType"`
	Digest      string            `json:"digest"`
	Annotations map[string]string `json:"annotations"`
	Platform    *struct {
		Architecture string `json:"architecture"`
		OS           string `json:"os"`
	} `json:"platform"`
}

type ociManifest struct {
	MediaType string          `json:"mediaType"`
	Manifests []ociDescriptor `json:"manifests"` // an index
	Config    ociDescriptor   `json:"config"`    // an image manifest
	Layers    []ociDescriptor `json:"layers"`
}

const refNameAnnotation = "org.opencontainers.image.ref.name"

func readOCILayout(dir string) ([]archiveImage, error) {
	raw, err := os.ReadFile(filepath.Join(dir, "index.json"))
	if err != nil {
		return nil, err
	}
	var index ociManifest
	if err := json.Unmarshal(raw, &index); err != nil {
		return nil, fmt.Errorf("index.json: %w", err)
	}
	var images []archiveImage
	for _, desc := range index.Manifests {
		img, err := resolveOCIDescriptor(dir, desc)
		if err != nil {
			return nil, err
		}
		if name := desc.Annotations[refNameAnnotation]; name != "" {
			img.names = []string{name}
		}
		images = append(images, img)
	}
	return images, nil
}

// resolveOCIDescriptor follows a descriptor to an image manifest,
// choosing the host's platform from a nested index.
func resolveOCIDescriptor(dir string, desc ociDescriptor) (archiveImage, error) {
	blob := blobPath(desc.Digest)
	raw, err := os.ReadFile(filepath.Join(dir, blob))
	if err != nil {
		return archiveImage{}, err
	}
	var m ociManifest
	if err := json.Unmarshal(raw, &m); err != nil {
		return archiveImage{}, fmt.Errorf("%s: %w", blob, err)
	}
	if len(m.Manifests) > 0 {
		var platforms []string
		for _, sub := range m.Manifests {
			if sub.Platform == nil {
				continue
			}
			platforms = append(platforms, sub.Platform.OS+"/"+sub.Platform.Architecture)
			if sub.Platform.OS == runtime.GOOS && sub.Platform.Architecture == runtime.GOARCH {
				return resolveOCIDescriptor(dir, sub)
			}
		}
		sort.Strings(platforms)
		return archiveImage{}, fmt.Errorf("%s: no manifest for %s/%s (available: %s)", blob, runtime.GOOS, runtime.GOARCH, strings.Join(platforms, ", "))
	}
	img := archiveImage{configPath: blobPath(m.Config.Digest)}
	for _, l := range m.Layers {
		img.layers = append(img.layers, blobPath(l.Digest))
	}
	return img, img.readConfig(dir)
}

func blobPath(digest string) string {
	alg, hex, _ := strings.Cut(digest, ":")
	return filepath.Join("blobs", alg, hex)
}

func (img *archiveImage) readConfig(dir string) error {
	if img.configPath == "" {
		return fmt.Errorf("image without a config blob")
	}
	raw, err := os.ReadFile(filepath.Join(dir, img.configPath))
	if err != nil {
		return err
	}
	if err := json.Unmarshal(raw, &img.config); err != nil {
		return fmt.Errorf("%s: %w", img.configPath, err)
	}
	return nil
}

// flattenImage applies the layers in order into dir/rootfs.
func flattenImage(unpacked string, img archiveImage, dir, archive string) error {
	rootfs := filepath.Join(dir, "rootfs")
	if err := os.MkdirAll(rootfs, 0o755); err != nil {
		return err
	}
	for _, layer := range img.layers {
		if err := applyLayer(rootfs, filepath.Join(unpacked, layer)); err != nil {
			return fmt.Errorf("%s: layer %s: %w", archive, layer, err)
		}
	}
	return nil
}

func writeMeta(dir string, meta ImageMeta) error {
	raw, err := json.MarshalIndent(meta, "", "  ")
	if err != nil {
		return err
	}
	return os.WriteFile(filepath.Join(dir, "image.json"), append(raw, '\n'), 0o644)
}

// place moves a finished image directory to its store location.
func (s Store) place(src string, ref Reference) error {
	dir := s.imageDir(ref)
	if err := os.MkdirAll(filepath.Dir(dir), 0o755); err != nil {
		return err
	}
	if err := os.Rename(src, dir); err != nil {
		s.pruneEmptyParents(dir)
		return err
	}
	return nil
}

// copyTree duplicates an image directory with hardlinks for files, so a
// second tag costs no data.
func copyTree(src, dst string) error {
	return filepath.Walk(src, func(p string, fi os.FileInfo, err error) error {
		if err != nil {
			return err
		}
		rel, _ := filepath.Rel(src, p)
		target := filepath.Join(dst, rel)
		switch {
		case fi.IsDir():
			return os.MkdirAll(target, fi.Mode()|0o700)
		case fi.Mode()&os.ModeSymlink != 0:
			link, err := os.Readlink(p)
			if err != nil {
				return err
			}
			return os.Symlink(link, target)
		default:
			return os.Link(p, target)
		}
	})
}