package-updater/internal/master/store/services.go

package store

import (
	"context"
	"errors"
	"fmt"
	"time"

	"github.com/jackc/pgx/v5"
	"github.com/jackc/pgx/v5/pgconn"

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

type execer interface {
	Exec(ctx context.Context, sql string, args ...any) (pgconn.CommandTag, error)
}

// lockActiveAgent takes a share lock on the agent's row inside tx and fails
// with model.ErrRevoked if the agent is revoked or gone. RevokeAgent's
// UPDATE of the row waits for the lock, so an inventory write either commits
// before the revoke (whose service cleanup then sees and removes its rows)
// or runs after it and is refused. Without this, an in-flight write could
// bring a revoked agent's services back.
func lockActiveAgent(ctx context.Context, tx pgx.Tx, agentID int64) error {
	var one int
	err := tx.QueryRow(ctx,
		`SELECT 1 FROM agents WHERE id = $1 AND revoked_at IS NULL FOR SHARE`, agentID).Scan(&one)
	if errors.Is(err, pgx.ErrNoRows) {
		return model.ErrRevoked
	}
	if err != nil {
		return fmt.Errorf("lock agent: %w", err)
	}
	return nil
}

const upsertServiceSQL = `
INSERT INTO services (
	kind, location_key, agent_id, compose_project, compose_service, working_dir, config_files,
	built, swarm_cluster_id, swarm_service, container_name, image_repo, actual_ref, actual_digest,
	status, self, first_seen_at, updated_at)
VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13, $14, $15, $16, $17, $17)
ON CONFLICT (location_key) DO UPDATE SET
	agent_id      = EXCLUDED.agent_id,
	working_dir   = EXCLUDED.working_dir,
	config_files  = EXCLUDED.config_files,
	built         = EXCLUDED.built,
	image_repo    = EXCLUDED.image_repo,
	actual_ref    = EXCLUDED.actual_ref,
	actual_digest = EXCLUDED.actual_digest,
	status        = EXCLUDED.status,
	self          = EXCLUDED.self,
	updated_at    = EXCLUDED.updated_at,
	removed_at    = NULL,
	awaiting_inventory = false`

func upsertService(ctx context.Context, db execer, agentID int64, clusterID string, svc protocol.Service, at time.Time) error {
	image := imageref.Parse(svc.Image)
	clusterCol := ""
	if svc.Kind == protocol.KindSwarm {
		clusterCol = clusterID
	}
	_, err := db.Exec(ctx, upsertServiceSQL,
		string(svc.Kind), model.LocationKey(agentID, clusterID, svc.Key()), agentID,
		svc.ComposeProject, svc.ComposeService, svc.WorkingDir, svc.ConfigFiles, svc.Built,
		clusterCol, svc.SwarmService, svc.ContainerName,
		image.Repo, svc.Image, svc.Digest,
		svc.Status, svc.Self, at)
	if err != nil {
		return fmt.Errorf("upsert service %s: %w", svc.Key(), err)
	}
	if imageref.IsDigest(svc.Digest) {
		// A bare image runs "latest", as jobs.FollowedRef sees it, so it
		// shares its row with the polled one. repo@digest has no tag.
		if image.Tag == "" && image.Digest == "" {
			image.Tag = "latest"
		}
		if err := insertVersion(ctx, db, image.Repo, image.Tag, svc.Digest, model.SourceRunning, at); err != nil {
			return err
		}
	}
	return nil
}

// insertVersion records a version in image_versions unless the same
// (repo, tag, digest) is already there; the first sighting keeps its source
// and first_seen_at.
func insertVersion(ctx context.Context, db execer, repo, tag, digest, source string, at time.Time) error {
	if _, err := db.Exec(ctx, `
		INSERT INTO image_versions (repo, tag, digest, source, first_seen_at)
		VALUES ($1, $2, $3, $4, $5) ON CONFLICT DO NOTHING`,
		repo, tag, digest, source, at); err != nil {
		return fmt.Errorf("record version %s: %w", imageref.Ref{Repo: repo, Tag: tag, Digest: digest}, err)
	}
	return nil
}

// ApplySnapshot makes the stored inventory for this agent's scope equal to
// services: everything listed is upserted, everything in scope but not
// listed is marked removed. Returns model.ErrRevoked without writing if the
// agent is revoked. Scope is the agent's own Compose services and
// containers, plus the cluster's Swarm services if the agent is a manager.
// A manager without a cluster ID is treated as a non-manager, since its
// Swarm rows could not be told apart from other agents' rows.
func (s *Store) ApplySnapshot(ctx context.Context, agentID int64, h protocol.Hello, services []protocol.Service, at time.Time) error {
	tx, err := s.pool.Begin(ctx)
	if err != nil {
		return fmt.Errorf("begin snapshot: %w", err)
	}
	defer func() { _ = tx.Rollback(ctx) }()
	if err := lockActiveAgent(ctx, tx, agentID); err != nil {
		return err
	}

	isManager := h.SwarmManager && h.SwarmClusterID != ""
	keys := make([]string, 0, len(services))
	for _, svc := range services {
		if svc.Kind == protocol.KindSwarm && !isManager {
			continue
		}
		if err := upsertService(ctx, tx, agentID, h.SwarmClusterID, svc, at); err != nil {
			return err
		}
		keys = append(keys, model.LocationKey(agentID, h.SwarmClusterID, svc.Key()))
	}
	if _, err := tx.Exec(ctx, `
		UPDATE services SET removed_at = $1
		WHERE removed_at IS NULL
		  AND NOT (location_key = ANY($2))
		  AND ((kind IN ('compose', 'container') AND agent_id = $3)
		    OR (kind = 'swarm' AND $4 AND swarm_cluster_id = $5))`,
		at, keys, agentID, isManager, h.SwarmClusterID); err != nil {
		return fmt.Errorf("mark removed services: %w", err)
	}
	if err := tx.Commit(ctx); err != nil {
		return fmt.Errorf("commit snapshot: %w", err)
	}
	return nil
}

// UpsertService stores one changed or new service. Returns
// model.ErrRevoked without writing if the agent is revoked.
func (s *Store) UpsertService(ctx context.Context, agentID int64, clusterID string, svc protocol.Service, at time.Time) error {
	tx, err := s.pool.Begin(ctx)
	if err != nil {
		return fmt.Errorf("begin upsert: %w", err)
	}
	defer func() { _ = tx.Rollback(ctx) }()
	if err := lockActiveAgent(ctx, tx, agentID); err != nil {
		return err
	}
	if err := upsertService(ctx, tx, agentID, clusterID, svc, at); err != nil {
		return err
	}
	if err := tx.Commit(ctx); err != nil {
		return fmt.Errorf("commit upsert: %w", err)
	}
	return nil
}

// RemoveService marks the service with the agent-relative key as removed.
func (s *Store) RemoveService(ctx context.Context, agentID int64, clusterID, agentKey string, at time.Time) error {
	if _, err := s.pool.Exec(ctx,
		`UPDATE services SET removed_at = $2 WHERE location_key = $1 AND removed_at IS NULL`,
		model.LocationKey(agentID, clusterID, agentKey), at); err != nil {
		return fmt.Errorf("remove service: %w", err)
	}
	return nil
}

const serviceColumns = `s.id, s.kind, s.location_key, s.agent_id, COALESCE(a.name, ''),
	s.compose_project, s.compose_service, s.working_dir, s.config_files, s.built, s.swarm_cluster_id,
	s.swarm_service, s.container_name, s.image_repo, s.desired_ref, s.actual_ref, s.actual_digest,
	s.status, s.self, s.update_mode, s.available_ref, s.available_digest, s.available_source, s.available_at,
	s.last_poll_error, s.last_tag_digest, s.awaiting_inventory, s.first_seen_at, s.updated_at, s.removed_at`

func scanService(row pgx.Row) (model.Service, error) {
	var m model.Service
	var kind string
	err := row.Scan(&m.ID, &kind, &m.LocationKey, &m.AgentID, &m.AgentName,
		&m.ComposeProject, &m.ComposeService, &m.WorkingDir, &m.ConfigFiles, &m.Built, &m.SwarmClusterID,
		&m.SwarmService, &m.ContainerName, &m.ImageRepo, &m.DesiredRef, &m.ActualRef, &m.ActualDigest,
		&m.Status, &m.Self, &m.UpdateMode, &m.AvailableRef, &m.AvailableDigest, &m.AvailableSource, &m.AvailableAt,
		&m.LastPollError, &m.LastTagDigest, &m.AwaitingInventory, &m.FirstSeenAt, &m.UpdatedAt, &m.RemovedAt)
	m.Kind = protocol.Kind(kind)
	return m, err
}

// ListServices returns all services not marked removed: machine-bound
// services grouped by agent name first, Swarm services last.
func (s *Store) ListServices(ctx context.Context) ([]model.Service, error) {
	return s.queryServices(ctx, "list services", "", nil)
}

// ServicesByRepo returns the services not marked removed whose image repo
// equals repo (case-sensitive), in the same order as ListServices.
func (s *Store) ServicesByRepo(ctx context.Context, repo string) ([]model.Service, error) {
	return s.queryServices(ctx, "services by repo", "AND s.image_repo = $1", []any{repo})
}

func (s *Store) queryServices(ctx context.Context, what, extraWhere string, args []any) ([]model.Service, error) {
	rows, err := s.pool.Query(ctx, `
		SELECT `+serviceColumns+`
		FROM services s LEFT JOIN agents a ON a.id = s.agent_id
		WHERE s.removed_at IS NULL `+extraWhere+`
		ORDER BY s.kind = 'swarm', COALESCE(a.name, ''), s.location_key`, args...)
	if err != nil {
		return nil, fmt.Errorf("%s: %w", what, err)
	}
	defer rows.Close()
	var out []model.Service
	for rows.Next() {
		m, err := scanService(rows)
		if err != nil {
			return nil, fmt.Errorf("scan service: %w", err)
		}
		out = append(out, m)
	}
	return out, rows.Err()
}

// SetAvailable records a newer ref for the service, replacing any earlier
// suggestion, and records the version in image_versions in the same
// transaction (the tag is parsed from a.Ref).
func (s *Store) SetAvailable(ctx context.Context, id int64, a model.Available) error {
	tx, err := s.pool.Begin(ctx)
	if err != nil {
		return fmt.Errorf("set available: begin: %w", err)
	}
	defer func() { _ = tx.Rollback(ctx) }()
	if _, err := tx.Exec(ctx, `
		UPDATE services SET available_ref = $2, available_digest = $3, available_source = $4, available_at = $5
		WHERE id = $1`, id, a.Ref, a.Digest, a.Source, a.At); err != nil {
		return fmt.Errorf("set available: %w", err)
	}
	ref := imageref.Parse(a.Ref)
	if err := insertVersion(ctx, tx, ref.Repo, ref.Tag, a.Digest, a.Source, a.At); err != nil {
		return err
	}
	if err := tx.Commit(ctx); err != nil {
		return fmt.Errorf("set available: commit: %w", err)
	}
	return nil
}

// SetPollError sets last_poll_error on the given services; msg "" clears it.
func (s *Store) SetPollError(ctx context.Context, ids []int64, msg string) error {
	if len(ids) == 0 {
		return nil
	}
	if _, err := s.pool.Exec(ctx,
		`UPDATE services SET last_poll_error = $2 WHERE id = ANY($1)`, ids, msg); err != nil {
		return fmt.Errorf("set poll error: %w", err)
	}
	return nil
}

// SetLastTagDigest sets last_tag_digest; see model.Service.LastTagDigest. The
// poller and jobs.RestoreDesired call it, and FinishJob sets the same column
// after a tag job. Inventory writes never change it.
func (s *Store) SetLastTagDigest(ctx context.Context, id int64, digest string) error {
	if _, err := s.pool.Exec(ctx,
		`UPDATE services SET last_tag_digest = $2 WHERE id = $1`, id, digest); err != nil {
		return fmt.Errorf("set last tag digest: %w", err)
	}
	return nil
}

// RecentVersions returns up to limit versions of repo, most recently first
// seen first. A version without a digest is left out when a version with the
// same tag and a digest was first seen at or after it.
func (s *Store) RecentVersions(ctx context.Context, repo string, limit int) ([]model.ImageVersion, error) {
	rows, err := s.pool.Query(ctx, `
		SELECT id, repo, tag, digest, source, first_seen_at FROM image_versions v
		WHERE v.repo = $1
		  AND NOT (v.digest = '' AND EXISTS (
		        SELECT 1 FROM image_versions d
		        WHERE d.repo = v.repo AND d.tag = v.tag AND d.digest <> ''
		          AND d.first_seen_at >= v.first_seen_at))
		ORDER BY v.first_seen_at DESC, v.id DESC
		LIMIT $2`, repo, limit)
	if err != nil {
		return nil, fmt.Errorf("recent versions: %w", err)
	}
	defer rows.Close()
	var out []model.ImageVersion
	for rows.Next() {
		var v model.ImageVersion
		if err := rows.Scan(&v.ID, &v.Repo, &v.Tag, &v.Digest, &v.Source, &v.FirstSeenAt); err != nil {
			return nil, fmt.Errorf("scan version: %w", err)
		}
		out = append(out, v)
	}
	return out, rows.Err()
}

// GetImageVersion returns one version. Returns model.ErrNotFound if the ID
// does not exist.
func (s *Store) GetImageVersion(ctx context.Context, id int64) (model.ImageVersion, error) {
	var v model.ImageVersion
	err := s.pool.QueryRow(ctx, `
		SELECT id, repo, tag, digest, source, first_seen_at FROM image_versions WHERE id = $1`, id,
	).Scan(&v.ID, &v.Repo, &v.Tag, &v.Digest, &v.Source, &v.FirstSeenAt)
	if errors.Is(err, pgx.ErrNoRows) {
		return model.ImageVersion{}, model.ErrNotFound
	}
	if err != nil {
		return model.ImageVersion{}, fmt.Errorf("get image version: %w", err)
	}
	return v, nil
}

// SetUpdateMode sets the service's update mode (model.ModeManual or
// model.ModeAuto). Returns model.ErrNotFound if the service does not exist
// or is marked removed.
func (s *Store) SetUpdateMode(ctx context.Context, id int64, mode string) error {
	tag, err := s.pool.Exec(ctx,
		`UPDATE services SET update_mode = $2 WHERE id = $1 AND removed_at IS NULL`, id, mode)
	if err != nil {
		return fmt.Errorf("set update mode: %w", err)
	}
	if tag.RowsAffected() == 0 {
		return model.ErrNotFound
	}
	return nil
}

// AcceptRunning makes the running version the desired one: desired_ref =
// desired, provided the service is live, has been updated before and still
// reports actualRef and actualDigest (what the caller saw when it decided).
// It starts no job and leaves last_tag_digest, the available suggestion and
// the update mode alone. Returns model.ErrNotFound for an unknown or removed
// service and model.ErrConflict when the service was never updated or the
// running version changed meanwhile.
func (s *Store) AcceptRunning(ctx context.Context, id int64, actualRef, actualDigest, desired string) error {
	tag, err := s.pool.Exec(ctx, `
		UPDATE services SET desired_ref = $4, awaiting_inventory = false
		WHERE id = $1 AND removed_at IS NULL AND desired_ref <> ''
		  AND actual_ref = $2 AND actual_digest = $3`, id, actualRef, actualDigest, desired)
	if err != nil {
		return fmt.Errorf("accept running: %w", err)
	}
	if tag.RowsAffected() == 1 {
		return nil
	}
	svc, err := s.GetService(ctx, id)
	if err != nil {
		return err
	}
	if svc.RemovedAt != nil {
		return model.ErrNotFound
	}
	return model.ErrConflict
}

// GetService returns one service, including a removed one (RemovedAt set).
// Returns model.ErrNotFound if the ID does not exist.
func (s *Store) GetService(ctx context.Context, id int64) (model.Service, error) {
	m, err := scanService(s.pool.QueryRow(ctx, `
		SELECT `+serviceColumns+`
		FROM services s LEFT JOIN agents a ON a.id = s.agent_id
		WHERE s.id = $1`, id))
	if errors.Is(err, pgx.ErrNoRows) {
		return model.Service{}, model.ErrNotFound
	}
	if err != nil {
		return model.Service{}, fmt.Errorf("get service: %w", err)
	}
	return m, nil
}