package server

import (
	"context"
	"io"
	"slices"
	"time"

	"github.com/tidbcloud/observability/vm-migration/internal/exitcode"
	"github.com/tidbcloud/observability/vm-migration/internal/printer"
)

// retryInterval is the fixed delay between retry attempts.
const retryInterval = time.Minute

// retryJob re-executes an inner job while an attempt ends in PartialFailure.
//
// Hard errors, cancellation, and context expiry return immediately. Each retry
// is announced on progress and separated by interval. Warnings accumulate
// across attempts because the inner loop resets its own warnings each pass.
type retryJob struct {
	inner    job
	retries  int
	progress io.Writer
	interval time.Duration
	warnings []string
}

func (j *retryJob) Kind() jobKind { return j.inner.Kind() }

// Run executes the inner job and returns the last attempt's error.
//
// Run is not safe for concurrent use; the dispatcher calls Warnings only
// after Run returns.
func (j *retryJob) Run(ctx context.Context) error {
	for attempt := 0; ; attempt++ {
		err := j.inner.Run(ctx)
		j.collectWarnings()
		if err == nil || ctx.Err() != nil || exitcode.Classify(err) != exitcode.Partial {
			return err
		}
		if attempt >= j.retries {
			return err
		}
		printer.Printf(j.progress, "retry %d/%d in %s: %s", attempt+1, j.retries, j.interval, err)
		select {
		case <-ctx.Done():
			return ctx.Err()
		case <-time.After(j.interval):
		}
	}
}

// collectWarnings merges the inner job's current warnings into the
// cross-attempt union.
func (j *retryJob) collectWarnings() {
	source, ok := j.inner.(warningSource)
	if !ok {
		return
	}
	for _, w := range source.Warnings() {
		if !slices.Contains(j.warnings, w) {
			j.warnings = append(j.warnings, w)
		}
	}
}

// Warnings returns the deduped union of the inner job's warnings across
// attempts.
func (j *retryJob) Warnings() []string { return j.warnings }

// wrapRetry returns job unchanged when retries <= 0, else a retryJob that
// re-executes it on PartialFailure with the given between-attempt delay.
//
// The decorator forwards Kind and Warnings only; jobs implementing
// artifactWriter must not be wrapped.
func wrapRetry(job job, retries int, progress io.Writer, interval time.Duration) job {
	if retries <= 0 {
		return job
	}
	return &retryJob{inner: job, retries: retries, progress: progress, interval: interval}
}