package agent

import (
	"context"

	"github.com/robfig/cron/v3"
)

type Job interface {
	// Run executes the cron task. It should cancel whatever it
	// is doing upon receiving a signal from ctx.Done.
	Run(context.Context)
}

// CronScheduler can run tasks on a cron schedule, with cancellation.
type CronScheduler struct {
	runner *cron.Cron
}

func NewCronScheduler() *CronScheduler {
	return &CronScheduler{runner: cron.New()}
}

func (s *CronScheduler) Run(ctx context.Context) {
	s.runner.Start()
	<-ctx.Done()
	stopCtx := s.runner.Stop()

	// Block until all jobs are stopped
	<-stopCtx.Done()
}

// bit of a hack so that we can pass a context into cron.Schedule
type jobWrapper struct {
	ctx     context.Context //nolint:containedctx
	runFunc func(context.Context)
}

func (j *jobWrapper) Run() {
	j.runFunc(j.ctx)
}

func (s *CronScheduler) Schedule(ctx context.Context, schedule cron.Schedule, job Job) {
	wrapper := &jobWrapper{
		ctx:     ctx,
		runFunc: job.Run,
	}

	s.runner.Schedule(schedule, wrapper)
}
