// Package wait implements the retry/timeout loop that polls a database until a
// liveness attempt succeeds. The clock and the attempt are injected so the loop
// is fully testable without a real database or real time.
package wait

import (
	"context"
	"time"
)

// Clock abstracts time so tests can advance it instantly.
type Clock interface {
	Now() time.Time
	// Sleep blocks for d or until ctx is done, returning ctx.Err() if it was
	// cancelled first.
	Sleep(ctx context.Context, d time.Duration) error
}

// AttemptFunc performs one connect+ping. The ctx carries the per-attempt
// connect timeout.
type AttemptFunc func(ctx context.Context) error

// OnFailure is invoked after each failed attempt that will be retried.
type OnFailure func(attempt int, err error, elapsed, timeout, retryIn time.Duration)

// Params configures the loop.
type Params struct {
	Timeout        time.Duration // 0 = wait forever
	Interval       time.Duration
	ConnectTimeout time.Duration
}

// Result reports how a run ended.
type Result struct {
	OK        bool
	TimedOut  bool
	Cancelled bool
	Attempts  int
	Elapsed   time.Duration
	LastErr   error
}

// Run polls attempt until it succeeds, the timeout is exhausted, or ctx is
// cancelled. It always makes at least one attempt.
func Run(ctx context.Context, p Params, attempt AttemptFunc, onFail OnFailure, clk Clock) Result {
	start := clk.Now()
	var deadline time.Time
	if p.Timeout > 0 {
		deadline = start.Add(p.Timeout)
	}

	res := Result{}
	for {
		res.Attempts++

		actx, cancel := context.WithTimeout(ctx, p.ConnectTimeout)
		err := attempt(actx)
		cancel()

		now := clk.Now()
		res.Elapsed = now.Sub(start)

		if err == nil {
			res.OK = true
			return res
		}
		res.LastErr = err

		// A cancelled parent context (SIGTERM/SIGINT) wins over a retry.
		if ctx.Err() != nil {
			res.Cancelled = true
			return res
		}

		// No time budget left for another attempt.
		if p.Timeout > 0 && !now.Before(deadline) {
			res.TimedOut = true
			return res
		}

		sleep := p.Interval
		if p.Timeout > 0 {
			if remaining := deadline.Sub(now); remaining < sleep {
				sleep = remaining
			}
		}
		onFail(res.Attempts, err, res.Elapsed, p.Timeout, sleep)

		if serr := clk.Sleep(ctx, sleep); serr != nil {
			res.Cancelled = true
			return res
		}
	}
}

// RealClock is the production Clock backed by the wall clock.
type RealClock struct{}

func (RealClock) Now() time.Time { return time.Now() }

func (RealClock) Sleep(ctx context.Context, d time.Duration) error {
	if d <= 0 {
		return ctx.Err()
	}
	t := time.NewTimer(d)
	defer t.Stop()
	select {
	case <-ctx.Done():
		return ctx.Err()
	case <-t.C:
		return nil
	}
}