package main

// Removing the images no declaration uses, keeping the one each module runs and the one before
// (novox/hq ADR 0251 §5, to-be 51 "A machine's images").
//
// The weekly prune takes dangling images only, and an image a machine pulled by digest is not
// dangling: every version of every module a machine ever ran stays on it. This takes them, and
// keeps four things:
//
//   - every image a container on this machine uses, in any state;
//   - every image this machine's declaration names — what the mesh would send it now and what it was
//     last sent — as the controller's `images` verb answers. **No answer, no removal**: what the
//     declaration names is the one thing that cannot be guessed;
//   - every image younger than the floor (seven days, the weekly prune's week);
//   - in each line of images holding one of those, the newest image besides them: the previous
//     version, so going back needs no pull.
//
// A line is the images sharing a repository name, joined across names when one image carries more
// than one (a build's local tag and the store's name for the same image).
//
// Removal is by every name an image carries, one image at a time, never forced: the runtime itself refuses an image a container
// uses, and a refusal is reported for that image while the rest go on.

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"sort"
	"strings"
	"time"
)

// Asker asks a seat's verb or a module's tool on the mesh (the SDK's stdio.Ask).
type Asker func(key string, body any) (json.RawMessage, error)

// PruneImagesAsk is what docker_prune_images is asked.
type PruneImagesAsk struct {
	DryRun        bool
	Why           string
	OlderThanDays int
	Match         string
	Limit         int
}

// declaredImages is the controller's `images` answer (novox/hq ADR 0251 §5).
type declaredImages struct {
	Node   string `json:"node"`
	Images []struct {
		Image     string   `json:"image"`
		Resources []string `json:"resources"`
		In        []string `json:"in"`
	} `json:"images"`
	SentKnown bool `json:"sent_known"`
}

// imageInspected is the part of `docker image inspect` this reads.
type imageInspected struct {
	ID          string   `json:"Id"`
	RepoTags    []string `json:"RepoTags"`
	RepoDigests []string `json:"RepoDigests"`
	Created     string   `json:"Created"`
	Size        int64    `json:"Size"`
}

// PrunedImage is one image as the answer says it.
type PrunedImage struct {
	ID       string   `json:"id"`
	Names    []string `json:"names"`
	Created  string   `json:"created"`
	Size     int64    `json:"size_bytes"`
	Kept     bool     `json:"kept"`
	Why      []string `json:"why,omitempty"`
	Line     string   `json:"line"`
	Declared []string `json:"declared_by,omitempty"`
	UsedBy   []string `json:"used_by,omitempty"`
	created  time.Time
	fullID   string
}

const (
	keptUsed     = "used-by-container"
	keptDeclared = "declared"
	keptPrevious = "previous"
	keptYoung    = "younger-than-floor"
	keptUnknown  = "declaration-unknown"
)

// normalRef is an image reference as the runtime reports it: the default registry and its library
// namespace left out, and a bare name tagged latest.
func normalRef(ref string) string {
	ref = strings.TrimSpace(ref)
	for _, p := range []string{"docker.io/", "index.docker.io/", "registry-1.docker.io/"} {
		ref = strings.TrimPrefix(ref, p)
	}
	ref = strings.TrimPrefix(ref, "library/")
	if !strings.Contains(ref, "@") {
		name, tag := splitTag(ref)
		if tag == "" {
			ref = name + ":latest"
		}
	}
	return ref
}

// splitTag splits name:tag, minding a registry's port.
func splitTag(ref string) (string, string) {
	i := strings.LastIndex(ref, ":")
	if i < 0 || strings.Contains(ref[i:], "/") {
		return ref, ""
	}
	return ref[:i], ref[i+1:]
}

// repositoryOf is a reference without its tag or digest.
func repositoryOf(ref string) string {
	ref = normalRef(ref)
	if name, _, ok := strings.Cut(ref, "@"); ok {
		return name
	}
	name, _ := splitTag(ref)
	return name
}

func digestOf(ref string) string {
	if _, d, ok := strings.Cut(ref, "@"); ok {
		return d
	}
	return ""
}

// askDeclared asks the controller what this machine's declarations name. An error means unknown.
func (c *Client) askDeclared() (*declaredImages, error) {
	if c.Node == "" {
		return nil, errors.New("the runtime did not say which machine this is (MESH_NODE is empty)")
	}
	if c.Ask == nil {
		return nil, errors.New("this bundle cannot ask the mesh")
	}
	raw, err := c.Ask("seat:mesh-controller.images", map[string]any{"node": c.Node})
	if err != nil {
		return nil, fmt.Errorf("the controller's images verb did not answer: %w", err)
	}
	answer, output, ok := answerOf(raw)
	if !ok {
		return nil, fmt.Errorf("the controller refused images: %s", firstLine(output))
	}
	var d declaredImages
	if err := json.Unmarshal(answer, &d); err != nil {
		return nil, fmt.Errorf("the controller's images answer is not readable: %v", err)
	}
	if d.Node != "" && d.Node != c.Node {
		return nil, fmt.Errorf("the controller answered for %q, not for this machine %q", d.Node, c.Node)
	}
	return &d, nil
}

// answerOf reads a verb's answer whatever wraps it: the controller's {output, ok, answer}, a text the
// runtime handed over, or the protocol's content list (as mesh-delivery reads it).
func answerOf(raw json.RawMessage) (json.RawMessage, string, bool) {
	var s string
	if json.Unmarshal(raw, &s) == nil {
		return answerOf(json.RawMessage(s))
	}
	var m map[string]json.RawMessage
	if json.Unmarshal(raw, &m) != nil {
		return raw, string(raw), true
	}
	if content, has := m["content"]; has {
		var items []struct {
			Text string `json:"text"`
		}
		var e struct {
			IsError bool `json:"isError"`
		}
		_ = json.Unmarshal(raw, &e)
		if json.Unmarshal(content, &items) == nil && len(items) > 0 {
			inner, out, ok := answerOf(json.RawMessage(items[0].Text))
			return inner, out, ok && !e.IsError
		}
	}
	if okRaw, has := m["ok"]; has {
		var ok bool
		_ = json.Unmarshal(okRaw, &ok)
		var output string
		_ = json.Unmarshal(m["output"], &output)
		if answer, has := m["answer"]; has && ok {
			return answer, output, true
		}
		return nil, output, ok
	}
	return raw, string(raw), true
}

// images is every image that has a name, inspected. Dangling ones are the weekly prune's.
func (c *Client) namedImages(ctx context.Context) ([]imageInspected, error) {
	out, err := c.docker(ctx, "image", "ls", "--quiet", "--no-trunc")
	if err != nil {
		return nil, err
	}
	seen, ids := map[string]bool{}, []string{}
	for _, id := range lines(out) {
		if !seen[id] {
			seen[id] = true
			ids = append(ids, id)
		}
	}
	if len(ids) == 0 {
		return nil, nil
	}
	out, err = c.docker(ctx, append([]string{"image", "inspect"}, ids...)...)
	if err != nil {
		if strings.Contains(err.Error(), "No such") {
			return c.namedImages(ctx) // one went between the two calls: ask once more
		}
		return nil, err
	}
	var got []imageInspected
	if err := json.Unmarshal([]byte(out), &got); err != nil {
		return nil, fmt.Errorf("docker image inspect answered something that is not JSON: %v", err)
	}
	named := got[:0]
	for _, i := range got {
		if len(i.RepoTags)+len(i.RepoDigests) > 0 {
			named = append(named, i)
		}
	}
	return named, nil
}

// PruneImages removes, or with DryRun only lists, the images no container and no declaration uses.
func (c *Client) PruneImages(ctx context.Context, a PruneImagesAsk) (map[string]any, error) {
	if !a.DryRun && strings.TrimSpace(a.Why) == "" {
		return nil, errors.New("a real run needs why: nothing was removed")
	}
	if a.OlderThanDays < 0 {
		return nil, errors.New("older_than_days is at least 0")
	}
	if a.Limit <= 0 {
		a.Limit = 100
	}
	raw, err := c.namedImages(ctx)
	if err != nil {
		return nil, err
	}
	containers, err := c.inspectAll(ctx)
	if err != nil {
		return nil, err
	}
	users := map[string][]string{}
	for _, ct := range containers {
		users[ct.Image] = append(users[ct.Image], strings.TrimPrefix(ct.Name, "/"))
	}
	declared, askErr := c.askDeclared()

	// What the declarations name, by full reference and by digest.
	byRef, byDigest := map[string][]string{}, map[string][]string{}
	if declared != nil {
		for _, d := range declared.Images {
			who := strings.Join(d.Resources, ",")
			if who == "" {
				who = d.Image
			}
			ref := normalRef(d.Image)
			byRef[ref] = append(byRef[ref], who)
			if dg := digestOf(ref); dg != "" {
				byDigest[dg] = append(byDigest[dg], who)
			}
		}
	}

	now := c.Now()
	floor := now.Add(-time.Duration(a.OlderThanDays) * 24 * time.Hour)
	images := make([]*PrunedImage, 0, len(raw))
	parent := map[string]string{}
	var find func(string) string
	find = func(x string) string {
		if parent[x] == "" || parent[x] == x {
			parent[x] = x
			return x
		}
		parent[x] = find(parent[x])
		return parent[x]
	}
	for _, r := range raw {
		img := &PrunedImage{ID: shortID(r.ID), fullID: r.ID, Size: r.Size, Created: r.Created}
		img.created, _ = time.Parse(time.RFC3339Nano, r.Created)
		img.Names = append(append([]string{}, r.RepoTags...), r.RepoDigests...)
		img.UsedBy = users[r.ID]
		repos := []string{}
		for _, n := range img.Names {
			ref := normalRef(n)
			repos = append(repos, repositoryOf(ref))
			who := byRef[ref]
			if dg := digestOf(ref); dg != "" && len(who) == 0 {
				who = byDigest[dg]
			}
			img.Declared = append(img.Declared, who...)
		}
		sort.Strings(repos)
		for _, rp := range repos[1:] {
			parent[find(rp)] = find(repos[0])
		}
		find(repos[0])
		img.Line = repos[0]
		images = append(images, img)
	}
	for _, img := range images {
		img.Line = find(img.Line)
		if len(img.UsedBy) > 0 {
			img.Why = append(img.Why, keptUsed)
		}
		if len(img.Declared) > 0 {
			img.Why = append(img.Why, keptDeclared)
		}
	}
	// The previous version: in each line holding an image in use or declared, the newest one besides.
	lines := map[string][]*PrunedImage{}
	for _, img := range images {
		lines[img.Line] = append(lines[img.Line], img)
	}
	for _, members := range lines {
		current := false
		for _, m := range members {
			current = current || len(m.Why) > 0
		}
		if !current {
			continue
		}
		var newest *PrunedImage
		for _, m := range members {
			if len(m.Why) == 0 && (newest == nil || m.created.After(newest.created)) {
				newest = m
			}
		}
		if newest != nil {
			newest.Why = append(newest.Why, keptPrevious)
		}
	}
	for _, img := range images {
		if img.created.IsZero() || img.created.After(floor) {
			img.Why = append(img.Why, keptYoung)
		}
		if declared == nil {
			img.Why = append(img.Why, keptUnknown)
		}
		img.Kept = len(img.Why) > 0
	}

	counts := map[string]int{"images": len(images)}
	var candidates []*PrunedImage
	var candidateBytes int64
	for _, img := range images {
		for _, w := range img.Why {
			counts["kept_"+strings.ReplaceAll(w, "-", "_")]++
		}
		if img.Kept {
			counts["kept"]++
			continue
		}
		if a.Match != "" && !strings.Contains(img.Line+" "+strings.Join(img.Names, " "), a.Match) {
			continue
		}
		candidates = append(candidates, img)
		candidateBytes += img.Size
	}
	counts["candidates"] = len(candidates)

	answer := map[string]any{
		"dry_run": a.DryRun, "node": c.Node, "declaration_known": declared != nil,
		"older_than_days": a.OlderThanDays, "counts": counts,
		"bytes_candidate": candidateBytes,
		"note": "bytes_candidate sums each image's size, an upper bound: images share layers, so less is freed. " +
			"Dangling images are not taken here; the weekly prune takes them.",
	}
	if declared != nil {
		answer["sent_known"] = declared.SentKnown
	} else {
		answer["sent_known"] = false
		answer["declaration_unknown"] = askErr.Error()
	}
	if a.Match != "" {
		answer["match"] = a.Match
	}

	sort.Slice(images, func(i, j int) bool {
		if images[i].Kept != images[j].Kept {
			return !images[i].Kept
		}
		return images[i].Size > images[j].Size
	})
	shown := images
	if a.Match != "" {
		shown = shown[:0:0]
		for _, img := range images {
			if strings.Contains(img.Line+" "+strings.Join(img.Names, " "), a.Match) {
				shown = append(shown, img)
			}
		}
	}
	answer["shown"] = min(len(shown), a.Limit)
	answer["images"] = shown[:min(len(shown), a.Limit)]

	if a.DryRun {
		answer["said"] = fmt.Sprintf("dry run: %d of %d images would be removed (at most %s); nothing was removed. "+
			"Call again with dry_run false and why to remove them.", len(candidates), len(images), human(candidateBytes))
		return answer, nil
	}
	if declared == nil {
		answer["said"] = "nothing was removed: what this machine's declaration names is not known (" + askErr.Error() + ")"
		return answer, nil
	}
	removed, refused := []string{}, []map[string]string{}
	var freed int64
	for _, img := range candidates {
		// By every name it carries, never forced: removing an image's last name removes the image, and the
		// runtime refuses one a container uses. (By id, an image with two names needs force, which this
		// never uses.)
		args := []string{"image", "rm"}
		var bad error
		for _, n := range img.Names {
			ref, err := Ref(n)
			if err != nil {
				bad = err
				break
			}
			args = append(args, ref)
		}
		if bad != nil {
			refused = append(refused, map[string]string{"id": img.ID, "why": bad.Error()})
			continue
		}
		if _, err := c.docker(ctx, args...); err != nil {
			refused = append(refused, map[string]string{"id": img.ID, "why": err.Error()})
			continue
		}
		removed = append(removed, img.ID)
		freed += img.Size
	}
	answer["removed"], answer["refused"], answer["bytes_removed"] = removed, refused, freed
	answer["why"] = a.Why
	answer["said"] = fmt.Sprintf("removed %d image(s) (at most %s), %d refused by the runtime", len(removed), human(freed), len(refused))
	return answer, nil
}

func human(b int64) string {
	switch {
	case b >= 1<<30:
		return fmt.Sprintf("%.1f GiB", float64(b)/(1<<30))
	case b >= 1<<20:
		return fmt.Sprintf("%.1f MiB", float64(b)/(1<<20))
	}
	return fmt.Sprintf("%d B", b)
}