// Package poller checks the registry for new digests of the tags the
// services follow. A new digest becomes the service's available version, and
// a service in auto mode gets an update job. It also lists the tags of the
// repos whose followed tag looks like a version and suggests a newer one;
// that suggestion never starts a job.
package poller

import (
	"context"
	"errors"
	"log/slog"
	"slices"
	"strings"
	"sync"
	"time"

	"github.com/shcizo/package-updater/internal/imageref"
	"github.com/shcizo/package-updater/internal/master/events"
	"github.com/shcizo/package-updater/internal/master/jobs"
	"github.com/shcizo/package-updater/internal/master/model"
	"github.com/shcizo/package-updater/internal/master/tagver"
)

// lookupTimeout bounds one registry lookup.
const lookupTimeout = 30 * time.Second

// listTimeout bounds one tag listing, which may take several requests.
const listTimeout = 60 * time.Second

// maxErrorBytes caps the registry error stored in services.last_poll_error.
const maxErrorBytes = 500

// Store is the storage the poller needs. *store.Store satisfies it.
type Store interface {
	ListServices(ctx context.Context) ([]model.Service, error)
	GetService(ctx context.Context, id int64) (model.Service, error)
	SetAvailable(ctx context.Context, id int64, a model.Available) error
	SetLastTagDigest(ctx context.Context, id int64, digest string) error
	SetNewer(ctx context.Context, id int64, ref string, at time.Time) error
	SetPollError(ctx context.Context, ids []int64, msg string) error
}

// Jobs starts updates. *jobs.Service satisfies it.
type Jobs interface {
	Start(ctx context.Context, svc model.Service, req jobs.StartRequest) (model.Job, error)
}

// Registry looks up the digest a "repo:tag" ref points at and lists the tags
// of a repo. *registry.Client satisfies it.
type Registry interface {
	Digest(ctx context.Context, ref string) (string, error)
	Tags(ctx context.Context, repo string) ([]string, error)
}

// Publisher receives change signals. *events.Broker satisfies it.
type Publisher interface{ Publish(t events.Topic) }

// Config holds the poller settings.
type Config struct {
	Interval   time.Duration // time between the end of a cycle and the start of the next (POLL_INTERVAL)
	FirstDelay time.Duration // wait before the first cycle; 1 minute in production
	Workers    int           // registry requests in flight at once; 4 in production
}

// Poller runs polling cycles.
type Poller struct {
	store Store
	jobs  Jobs
	reg   Registry
	pub   Publisher
	cfg   Config
	log   *slog.Logger
	now   func() time.Time
}

// New returns a Poller.
func New(st Store, j Jobs, reg Registry, pub Publisher, cfg Config, log *slog.Logger) *Poller {
	if cfg.Workers < 1 {
		cfg.Workers = 1
	}
	return &Poller{store: st, jobs: j, reg: reg, pub: pub, cfg: cfg, log: log, now: time.Now}
}

// SkipReason says why svc is not polled: "follows no tag", "built from
// source", "no registry digest", or "" when it is polled.
func SkipReason(svc model.Service) string {
	switch {
	case jobs.FollowedRef(svc.ActualRef) == "":
		return model.ReasonNoTag
	case svc.Built:
		return "built from source"
	case svc.ActualDigest == "":
		return "no registry digest"
	}
	return ""
}

// Run polls until ctx is done. The first cycle starts Config.FirstDelay after
// the call and each following one Config.Interval after the previous cycle
// ended, so cycles never overlap. It blocks.
func (p *Poller) Run(ctx context.Context) {
	timer := time.NewTimer(p.cfg.FirstDelay)
	defer timer.Stop()
	for {
		select {
		case <-ctx.Done():
			return
		case <-timer.C:
		}
		p.Cycle(ctx)
		timer.Reset(p.cfg.Interval)
	}
}

// lookup is the registry's answer for one followed ref.
type lookup struct {
	ref    string
	digest string
	err    error
}

// cycleStats is what one cycle did, for the poll_cycle log line.
type cycleStats struct {
	refs, ok, failed, suggested, jobs, skipped int
	reposListed, tagsFailed, newer             int
	wrote                                      bool
}

// Cycle runs one poll over all services. It looks up every distinct followed
// ref once, with at most Config.Workers lookups at a time, then records the
// results. Then it lists the tags of each repo that has a service following a
// version-like tag, once per repo, and records newer tags. Writing results and
// starting jobs is not cut short by ctx ending during the lookups, so a job
// that was marked is always started.
func (p *Poller) Cycle(ctx context.Context) {
	began := p.now()
	svcs, err := p.store.ListServices(ctx)
	if err != nil {
		p.log.Error("poll_list_failed", "error", err)
		return
	}
	var stats cycleStats
	byRef := map[string][]model.Service{}
	var refs []string
	for _, svc := range svcs {
		if svc.RemovedAt != nil {
			continue
		}
		if SkipReason(svc) != "" {
			stats.skipped++
			continue
		}
		ref := jobs.FollowedRef(svc.ActualRef)
		if _, seen := byRef[ref]; !seen {
			refs = append(refs, ref)
		}
		byRef[ref] = append(byRef[ref], svc)
	}
	stats.refs = len(refs)

	write := context.WithoutCancel(ctx)
	for _, l := range p.lookupAll(ctx, refs) {
		if l.err != nil && ctx.Err() != nil {
			continue // cancelled mid-lookup: not the registry's fault
		}
		p.record(write, l, byRef[l.ref], &stats)
	}
	// Publish after the digest phase too: listing tags can take a while and
	// an "Update available" should not wait for it.
	p.publishIfWrote(&stats)
	p.syncNewer(ctx, write, svcs, &stats)
	p.publishIfWrote(&stats)
	p.log.Info("poll_cycle", "refs", stats.refs, "ok", stats.ok, "failed", stats.failed,
		"suggested", stats.suggested, "jobs", stats.jobs, "skipped", stats.skipped,
		"repos_listed", stats.reposListed, "tags_failed", stats.tagsFailed, "newer", stats.newer,
		"duration_ms", p.now().Sub(began).Milliseconds())
}

// publishIfWrote tells the browsers to refresh if a phase of the cycle wrote
// anything, and starts the next phase with a clean slate.
func (p *Poller) publishIfWrote(stats *cycleStats) {
	if stats.wrote {
		p.pub.Publish(events.Services)
		stats.wrote = false
	}
}

// forEach calls fn(i) for i in [0, n), Config.Workers at a time, and stops
// starting new calls once ctx is done. It returns when all started calls
// have returned.
func (p *Poller) forEach(ctx context.Context, n int, fn func(i int)) {
	next := make(chan int)
	var wg sync.WaitGroup
	for w := 0; w < p.cfg.Workers; w++ {
		wg.Add(1)
		go func() {
			defer wg.Done()
			for i := range next {
				fn(i)
			}
		}()
	}
feed:
	for i := 0; i < n; i++ {
		select {
		case next <- i:
		case <-ctx.Done():
			break feed
		}
	}
	close(next)
	wg.Wait()
}

// lookupAll asks the registry for every ref, Config.Workers at a time. It
// returns the answers in ref order and stops starting new lookups once ctx
// is done; refs never looked up are missing from the result.
func (p *Poller) lookupAll(ctx context.Context, refs []string) []lookup {
	results := make([]*lookup, len(refs))
	p.forEach(ctx, len(refs), func(i int) {
		lctx, cancel := context.WithTimeout(ctx, lookupTimeout)
		d, err := p.reg.Digest(lctx, refs[i])
		cancel()
		results[i] = &lookup{ref: refs[i], digest: d, err: err}
	})
	var out []lookup
	for _, r := range results {
		if r != nil {
			out = append(out, *r)
		}
	}
	return out
}

// tagList is the registry's answer for the tags of one repo.
type tagList struct {
	repo string
	tags []string
	err  error
}

// listAll lists the tags of every repo like lookupAll looks up refs.
func (p *Poller) listAll(ctx context.Context, repos []string) []tagList {
	results := make([]*tagList, len(repos))
	p.forEach(ctx, len(repos), func(i int) {
		lctx, cancel := context.WithTimeout(ctx, listTimeout)
		tags, err := p.reg.Tags(lctx, repos[i])
		cancel()
		results[i] = &tagList{repo: repos[i], tags: tags, err: err}
	})
	var out []tagList
	for _, r := range results {
		if r != nil {
			out = append(out, *r)
		}
	}
	return out
}

// Followed returns the repo and tag svc follows: a bare image follows
// "latest", and a digest-only reference follows no tag.
func Followed(svc model.Service) (repo, tag string) {
	ref := imageref.Parse(jobs.FollowedRef(svc.ActualRef))
	return ref.Repo, ref.Tag
}

// followedComparable returns the repo and tag svc follows. ok is false unless
// svc is polled and the tag can have a newer tag.
func followedComparable(svc model.Service) (repo, tag string, ok bool) {
	if SkipReason(svc) != "" {
		return "", "", false
	}
	repo, tag = Followed(svc)
	return repo, tag, tagver.Comparable(tag)
}

// syncNewer lists the tags of every repo that has a service following a
// version-like tag, once per repo, and stores the newer tag of each service.
// A service whose tag cannot have a newer tag any more gets its suggestion
// cleared; one whose repo could not be listed keeps it, unless it names
// another repo than the service follows now. Nothing is written
// when the value is unchanged. It never starts a job, not even for an auto
// service: a newer tag may be a major step the user has not chosen.
func (p *Poller) syncNewer(ctx, write context.Context, svcs []model.Service, stats *cycleStats) {
	seen := map[string]bool{}
	var repos []string
	for _, svc := range svcs {
		if repo, _, ok := followedComparable(svc); ok && svc.RemovedAt == nil && !seen[repo] {
			seen[repo] = true
			repos = append(repos, repo)
		}
	}
	lists := map[string][]string{}
	for _, l := range p.listAll(ctx, repos) {
		if l.err != nil {
			p.log.Warn("poll_tags_failed", "repo", l.repo, "error", truncate(l.err.Error()))
			if ctx.Err() == nil { // cancelled mid-listing: not the registry's fault
				stats.tagsFailed++
			}
			continue
		}
		stats.reposListed++
		lists[l.repo] = l.tags
	}

	type repoTag struct{ repo, tag string }
	newest := map[repoTag]string{}
	for _, svc := range svcs {
		if svc.RemovedAt != nil {
			continue
		}
		want := ""
		if repo, tag, ok := followedComparable(svc); ok {
			tags, listed := lists[repo]
			if !listed && imageref.Parse(svc.NewerRef).Repo == repo {
				if svc.NewerRef != "" {
					stats.newer++
				}
				continue // not listed this cycle: keep what is stored
			}
			if listed {
				key := repoTag{repo, tag}
				n, done := newest[key]
				if !done {
					n, _ = tagver.Newer(tag, tags)
					newest[key] = n
				}
				if n != "" {
					want = repo + ":" + n
				}
			}
		}
		if want != "" {
			stats.newer++
		}
		if want == svc.NewerRef {
			continue
		}
		if err := p.store.SetNewer(write, svc.ID, want, p.now()); err != nil {
			p.log.Error("poll_set_newer_failed", "service_id", svc.ID, "ref", want, "error", err)
			continue
		}
		stats.wrote = true
	}
}

// record stores one lookup result for the services that follow l.ref.
func (p *Poller) record(ctx context.Context, l lookup, svcs []model.Service, stats *cycleStats) {
	if l.err != nil {
		stats.failed++
		msg := truncate(l.err.Error())
		p.log.Warn("poll_failed", "ref", l.ref, "error", msg)
		p.setPollError(ctx, svcs, msg, stats, func(s model.Service) bool { return s.LastPollError != msg })
		return
	}
	stats.ok++
	p.setPollError(ctx, svcs, "", stats, func(s model.Service) bool { return s.LastPollError != "" })
	var auto []model.Service
	for _, svc := range svcs {
		if !p.suggest(ctx, svc, l, stats) {
			continue
		}
		if autoWanted(svc, l.ref, l.digest) {
			auto = append(auto, svc)
		}
	}
	// An agent's own service goes last: jobs that reach an agent while it
	// replaces itself are lost, and messages to one agent arrive in order.
	slices.SortStableFunc(auto, func(a, b model.Service) int { return selfRank(a) - selfRank(b) })
	for _, svc := range auto {
		p.startAuto(ctx, svc.ID, l.ref, l.digest, stats)
	}
}

// selfRank sorts an agent's own service after all others.
func selfRank(svc model.Service) int {
	if svc.Self {
		return 1
	}
	return 0
}

// suggest marks l.digest as the available version of svc unless the service
// already runs it or already has exactly that suggestion. A service that
// runs l.digest also gets it as last_tag_digest. It returns false only when
// writing the suggestion failed.
func (p *Poller) suggest(ctx context.Context, svc model.Service, l lookup, stats *cycleStats) bool {
	if svc.ActualDigest == l.digest && svc.LastTagDigest != l.digest {
		// Seed last_tag_digest for a service that was not updated by a tag job
		// (it predates the column or was updated by hand), so that a later
		// rollback is not undone.
		if err := p.store.SetLastTagDigest(ctx, svc.ID, l.digest); err != nil {
			p.log.Error("poll_set_last_tag_digest_failed", "service_id", svc.ID, "error", err)
		}
	}
	if svc.ActualDigest == l.digest || (svc.AvailableRef == l.ref && svc.AvailableDigest == l.digest) {
		return true
	}
	err := p.store.SetAvailable(ctx, svc.ID, model.Available{Ref: l.ref, Digest: l.digest, Source: model.SourcePoll, At: p.now()})
	if err != nil {
		p.log.Error("poll_set_available_failed", "service_id", svc.ID, "ref", l.ref, "error", err)
		return false
	}
	stats.wrote = true
	stats.suggested++
	return true
}

// autoWanted reports whether an update of svc to ref at digest should be
// started automatically: the service is in auto mode, not removed, does not
// already run that digest, was not already updated automatically to it
// (otherwise a rollback to an older version would be undone at the next
// cycle), and its last job did not already deploy it (desired_ref holds the
// digest until the agent reports the new inventory).
func autoWanted(svc model.Service, ref, digest string) bool {
	return svc.RemovedAt == nil &&
		svc.UpdateMode == model.ModeAuto &&
		svc.ActualDigest != digest &&
		svc.LastTagDigest != digest &&
		svc.DesiredRef != ref+"@"+digest
}

// startAuto starts the update job of an auto service to ref and remembers
// digest as services.last_tag_digest, so the same digest is not started
// again, for instance after the user rolled back. It re-reads the service
// first, because a cycle can run for a long time and the user may have
// switched the service to manual meanwhile. A service that already has a
// job or cannot be updated is not an error for the cycle, and last_tag_digest
// is left alone so the next cycle tries again.
func (p *Poller) startAuto(ctx context.Context, id int64, ref, digest string, stats *cycleStats) {
	svc, err := p.store.GetService(ctx, id)
	if err != nil {
		p.log.Warn("poll_get_service_failed", "service_id", id, "error", err)
		return
	}
	if !autoWanted(svc, ref, digest) {
		return
	}
	_, err = p.jobs.Start(ctx, svc, jobs.StartRequest{Trigger: model.TriggerPoll, Ref: ref})
	var notUpdatable *model.NotUpdatableError
	switch {
	case err == nil:
		stats.jobs++
		if err := p.store.SetLastTagDigest(ctx, id, digest); err != nil {
			p.log.Error("poll_set_last_tag_digest_failed", "service_id", id, "error", err)
		}
	case errors.Is(err, model.ErrJobActive), errors.As(err, &notUpdatable):
		p.log.Debug("poll_job_not_started", "service_id", id, "ref", ref, "reason", err.Error())
	default:
		p.log.Warn("poll_job_failed", "service_id", id, "ref", ref, "error", err)
	}
}

// setPollError sets msg as last_poll_error on the services for which
// changed returns true; with none, it writes nothing.
func (p *Poller) setPollError(ctx context.Context, svcs []model.Service, msg string, stats *cycleStats, changed func(model.Service) bool) {
	var ids []int64
	for _, svc := range svcs {
		if changed(svc) {
			ids = append(ids, svc.ID)
		}
	}
	if len(ids) == 0 {
		return
	}
	if err := p.store.SetPollError(ctx, ids, msg); err != nil {
		p.log.Error("poll_set_error_failed", "error", err)
		return
	}
	stats.wrote = true
}

// truncate cuts msg to maxErrorBytes without leaving a broken UTF-8
// sequence, and drops NUL bytes, which Postgres text cannot hold.
func truncate(msg string) string {
	msg = strings.ReplaceAll(msg, "\x00", "")
	if len(msg) > maxErrorBytes {
		msg = strings.ToValidUTF8(msg[:maxErrorBytes], "")
	}
	return msg
}