package-updater/internal/agent/inventory/collector.go

// Package inventory turns what the local Docker daemon runs into the
// services an agent reports to the master.
package inventory

import (
	"context"
	"fmt"
	"log/slog"
	"os"
	"sort"
	"strings"
	"sync"
	"time"

	"github.com/docker/docker/api/types/container"
	"github.com/docker/docker/api/types/image"
	"github.com/docker/docker/api/types/swarm"
	"github.com/docker/docker/api/types/system"

	"github.com/shcizo/package-updater/internal/imageref"
	"github.com/shcizo/package-updater/internal/protocol"
)

const (
	labelProject     = "com.docker.compose.project"
	labelService     = "com.docker.compose.service"
	labelWorkingDir  = "com.docker.compose.project.working_dir"
	labelConfigFiles = "com.docker.compose.project.config_files"
	labelEnvFile     = "com.docker.compose.project.environment_file"
	labelSwarmTask   = "com.docker.swarm.service.id"
)

// DockerClient is the subset of the Docker SDK the collector needs.
// *client.Client satisfies it.
type DockerClient interface {
	ContainerList(ctx context.Context, opts container.ListOptions) ([]container.Summary, error)
	ContainerInspect(ctx context.Context, containerID string) (container.InspectResponse, error)
	ImageList(ctx context.Context, opts image.ListOptions) ([]image.Summary, error)
	ServiceList(ctx context.Context, opts swarm.ServiceListOptions) ([]swarm.Service, error)
	Info(ctx context.Context) (system.Info, error)
}

// BuildCheck finds the services of a Compose project that are built from
// source. The arguments are labels of one of the project's containers.
// *updater.BuildChecker satisfies it.
type BuildCheck interface {
	BuiltServices(ctx context.Context, project, workingDir, configFilesLabel, envFileLabel string) (map[string]bool, error)
}

// buildRetry is how long a failed build check is not repeated while the
// project's files stay the same.
const buildRetry = 5 * time.Minute

// Collector reads the local daemon's inventory.
type Collector struct {
	cli    DockerClient
	selfID string

	check BuildCheck
	log   *slog.Logger
	now   func() time.Time

	mu     sync.Mutex
	cache  map[string]buildEntry // by Compose project
	warned map[string]bool       // logged warnings: "<project>" (check failed) and "<project>/<service>" (service not in config)
}

// buildEntry is the cached outcome of a build check. It is valid as long as
// key (the project's labels plus its files' sizes and mtimes) is unchanged.
type buildEntry struct {
	key   string
	built map[string]bool
	err   error
	at    time.Time
}

// stackLabels are the Compose labels a build check needs.
type stackLabels struct{ workingDir, configFiles, envFile string }

// New returns a Collector. selfID is a prefix of the agent's own container
// ID (inside a container, the hostname); that container is marked Self.
func New(cli DockerClient, selfID string) *Collector {
	return &Collector{cli: cli, selfID: selfID, now: time.Now, cache: map[string]buildEntry{}, warned: map[string]bool{}}
}

// WithBuildCheck makes the collector mark Compose services that have a
// build: section as Built. The check re-runs only when a project's files
// change. When it fails, the project's services are reported as not built and
// the failure is logged once. It returns c.
func (c *Collector) WithBuildCheck(check BuildCheck, log *slog.Logger) *Collector {
	c.check, c.log = check, log
	return c
}

// Hello returns the machine facts sent on connect. AgentVersion is left for
// the caller to fill in.
func (c *Collector) Hello(ctx context.Context) (protocol.Hello, error) {
	info, err := c.cli.Info(ctx)
	if err != nil {
		return protocol.Hello{}, fmt.Errorf("docker info: %w", err)
	}
	h := protocol.Hello{Hostname: info.Name, SwarmManager: info.Swarm.ControlAvailable}
	if info.Swarm.Cluster != nil {
		h.SwarmClusterID = info.Swarm.Cluster.ID
	}
	return h, nil
}

// Snapshot returns every service on this machine, sorted by Key.
func (c *Collector) Snapshot(ctx context.Context) ([]protocol.Service, error) {
	info, err := c.cli.Info(ctx)
	if err != nil {
		return nil, fmt.Errorf("docker info: %w", err)
	}
	containers, err := c.cli.ContainerList(ctx, container.ListOptions{All: true})
	if err != nil {
		return nil, fmt.Errorf("docker container list: %w", err)
	}
	images, err := c.cli.ImageList(ctx, image.ListOptions{})
	if err != nil {
		return nil, fmt.Errorf("docker image list: %w", err)
	}
	repoDigests := make(map[string][]string, len(images))
	for _, img := range images {
		repoDigests[img.ID] = img.RepoDigests
	}

	byKey := make(map[string]protocol.Service)
	stacks := make(map[string]stackLabels)
	for _, ct := range containers {
		if _, isTask := ct.Labels[labelSwarmTask]; isTask {
			continue
		}
		ref, err := c.imageRef(ctx, ct)
		if err != nil {
			return nil, err
		}
		svc := c.fromContainer(ct, ref, repoDigests[ct.ImageID])
		if svc.Kind == protocol.KindCompose {
			if _, ok := stacks[svc.ComposeProject]; !ok {
				stacks[svc.ComposeProject] = stackLabels{svc.WorkingDir, svc.ConfigFiles, ct.Labels[labelEnvFile]}
			}
		}
		if prev, ok := byKey[svc.Key()]; ok {
			svc = mergeReplicas(prev, svc)
		}
		byKey[svc.Key()] = svc
	}

	c.markBuilt(ctx, byKey, stacks)

	if info.Swarm.ControlAvailable {
		services, err := c.cli.ServiceList(ctx, swarm.ServiceListOptions{Status: true})
		if err != nil {
			return nil, fmt.Errorf("docker service list: %w", err)
		}
		for _, s := range services {
			if s.Spec.TaskTemplate.ContainerSpec == nil {
				continue
			}
			svc := fromSwarmService(s)
			byKey[svc.Key()] = svc
		}
	}

	out := make([]protocol.Service, 0, len(byKey))
	for _, s := range byKey {
		out = append(out, s)
	}
	sort.Slice(out, func(i, j int) bool { return out[i].Key() < out[j].Key() })
	return out, nil
}

// imageRef returns the image reference the container was created from.
// When the tag has since moved to another image, ContainerList reports the
// image ID ("sha256:...") instead, so the reference is read from the
// container's config.
func (c *Collector) imageRef(ctx context.Context, ct container.Summary) (string, error) {
	if !strings.HasPrefix(ct.Image, "sha256:") {
		return ct.Image, nil
	}
	insp, err := c.cli.ContainerInspect(ctx, ct.ID)
	if err != nil {
		return "", fmt.Errorf("docker container inspect %s: %w", ct.ID, err)
	}
	if insp.Config == nil {
		return ct.Image, nil
	}
	return insp.Config.Image, nil
}

func (c *Collector) fromContainer(ct container.Summary, ref string, repoDigests []string) protocol.Service {
	svc := protocol.Service{
		Image:  ref,
		Digest: digestFor(ref, repoDigests),
		Status: containerStatus(ct.State),
		Self:   c.selfID != "" && strings.HasPrefix(ct.ID, c.selfID),
	}
	project, service := ct.Labels[labelProject], ct.Labels[labelService]
	if project != "" && service != "" {
		svc.Kind = protocol.KindCompose
		svc.ComposeProject = project
		svc.ComposeService = service
		svc.WorkingDir = ct.Labels[labelWorkingDir]
		svc.ConfigFiles = ct.Labels[labelConfigFiles]
		return svc
	}
	svc.Kind = protocol.KindContainer
	svc.ContainerName = ct.ID
	if len(ct.Names) > 0 {
		svc.ContainerName = strings.TrimPrefix(ct.Names[0], "/")
	}
	return svc
}

func fromSwarmService(s swarm.Service) protocol.Service {
	ref := s.Spec.TaskTemplate.ContainerSpec.Image
	return protocol.Service{
		Kind:         protocol.KindSwarm,
		SwarmService: s.Spec.Name,
		Image:        ref,
		Digest:       imageref.Parse(ref).Digest,
		Status:       swarmStatus(s.ServiceStatus),
	}
}

// mergeReplicas folds a second container of the same Compose service into
// the first: the service counts as running if any replica runs.
func mergeReplicas(a, b protocol.Service) protocol.Service {
	if b.Status == "running" {
		a.Status = "running"
	}
	a.Self = a.Self || b.Self
	return a
}

// digestFor returns the registry digest the container's image was pulled
// as. A reference that already pins a digest wins; otherwise the image's
// RepoDigests entry for the same repository is used.
func digestFor(ref string, repoDigests []string) string {
	r := imageref.Parse(ref)
	if r.Digest != "" {
		return r.Digest
	}
	for _, rd := range repoDigests {
		if p := imageref.Parse(rd); p.Repo == r.Repo {
			return p.Digest
		}
	}
	return ""
}

func containerStatus(state string) string {
	if state == "running" {
		return "running"
	}
	return "stopped"
}

func swarmStatus(st *swarm.ServiceStatus) string {
	switch {
	case st == nil:
		return "unknown"
	case st.DesiredTasks == 0:
		return "stopped"
	case st.RunningTasks >= st.DesiredTasks:
		return "running"
	default:
		return "degraded"
	}
}

// markBuilt sets Built on Compose services that have a build: section.
func (c *Collector) markBuilt(ctx context.Context, byKey map[string]protocol.Service, stacks map[string]stackLabels) {
	if c.check == nil {
		return
	}
	for key, svc := range byKey {
		if svc.Kind != protocol.KindCompose {
			continue
		}
		built := c.builtServices(ctx, svc.ComposeProject, stacks[svc.ComposeProject])
		if built == nil {
			continue // check failed; already logged
		}
		isBuilt, known := built[svc.ComposeService]
		if !known {
			// Fail closed: a service compose does not list cannot be shown to
			// be free of build:, so it is not offered for update.
			c.warnOnce(svc.ComposeProject+"/"+svc.ComposeService, "compose_service_not_in_config",
				"project", svc.ComposeProject, "service", svc.ComposeService)
		}
		if isBuilt || !known {
			svc.Built = true
			byKey[key] = svc
		}
	}
}

// builtServices returns the project's built services, from the cache while
// the project's files are unchanged.
func (c *Collector) builtServices(ctx context.Context, project string, l stackLabels) map[string]bool {
	// The key covers the files named in the labels. Changes in files they
	// pull in (extends:, include:, the default .env) do not invalidate the
	// cache, so Built can be stale until another change or an agent restart.
	// The job-time check is always fresh, so jobs stay safe.
	key := l.workingDir + "|" + l.configFiles + "|" + l.envFile + "|" + fileStamps(l.configFiles+","+l.envFile)
	c.mu.Lock()
	defer c.mu.Unlock()
	if e, ok := c.cache[project]; ok && e.key == key && (e.err == nil || c.now().Sub(e.at) < buildRetry) {
		return e.built
	}
	built, err := c.check.BuiltServices(ctx, project, l.workingDir, l.configFiles, l.envFile)
	c.cache[project] = buildEntry{key: key, built: built, err: err, at: c.now()}
	if err == nil {
		// Forget the project's warnings so they are logged again if the
		// problem returns. Project names cannot contain "/".
		for k := range c.warned {
			if k == project || strings.HasPrefix(k, project+"/") {
				delete(c.warned, k)
			}
		}
	} else {
		c.warnOnceLocked(project, "compose_build_check_failed", "project", project, "error", err)
	}
	return built
}

// warnOnce logs msg the first time it is called for key. key is a project
// name (its build check failed) or "<project>/<service>" (the service is not
// in the project's config); a successful check of the project forgets both.
func (c *Collector) warnOnce(key, msg string, args ...any) {
	c.mu.Lock()
	defer c.mu.Unlock()
	c.warnOnceLocked(key, msg, args...)
}

func (c *Collector) warnOnceLocked(key, msg string, args ...any) {
	if c.warned[key] || c.log == nil {
		return
	}
	c.warned[key] = true
	c.log.Warn(msg, args...)
}

// fileStamps summarises size and mtime of the comma-separated files, so a
// changed file changes the string.
func fileStamps(list string) string {
	var b strings.Builder
	for _, f := range strings.Split(list, ",") {
		if f = strings.TrimSpace(f); f == "" {
			continue
		}
		if fi, err := os.Stat(f); err == nil {
			fmt.Fprintf(&b, "%s:%d:%d;", f, fi.Size(), fi.ModTime().UnixNano())
		} else {
			fmt.Fprintf(&b, "%s:missing;", f)
		}
	}
	return b.String()
}