package util

import (
	"context"
	"errors"
	"fmt"
	"log/slog"
	"os"
	"sync"
	"testing"
	"time"
)

func TestMain(m *testing.M) {
	// Keep worker log lines out of the test output.
	slog.SetDefault(slog.New(slog.DiscardHandler))
	os.Exit(m.Run())
}

func TestWorkerRunRejectsBadConfig(t *testing.T) {
	noop := func(context.Context, time.Time) error {
		return nil
	}
	tests := []struct {
		name string
		w    Worker
	}{
		{"no task", Worker{
			Name:     "a",
			Interval: time.Second,
		}},
		{"zero interval", Worker{
			Name: "b",
			Task: noop,
		}},
		{"negative interval", Worker{
			Name:     "c",
			Interval: -time.Second,
			Task:     noop,
		}},
	}
	for _, tc := range tests {
		t.Run(tc.name, func(t *testing.T) {
			if err := tc.w.Run(context.Background(), time.Time{}); err == nil {
				t.Fatal("expected an error")
			}
		})
	}
}

func TestWorkerRunCancelledBeforeStart(t *testing.T) {
	ctx, cancel := context.WithCancel(context.Background())
	cancel()
	called := false
	w := Worker{
		Name:     "w",
		Interval: time.Millisecond,
		Task: func(context.Context, time.Time) error {
			called = true
			return nil
		},
	}

	if err := w.Run(ctx, time.Time{}); err != nil {
		t.Fatal(err)
	}
	if called {
		t.Fatal("task ran after the context was already cancelled")
	}
}

// The checkpoint passed to each run advances only when the previous run
// succeeded or reported ErrWorkerNoWork.
func TestWorkerRunCheckpoints(t *testing.T) {
	errFail := errors.New("boom")
	results := []error{
		nil,                                  // advance
		errFail,                              // keep
		context.Canceled,                     // keep
		fmt.Errorf("%w: x", ErrWorkerNoWork), // advance
		nil,
	}
	since0 := time.Date(2026, 1, 1, 0, 0, 0, 0, time.UTC)

	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()
	var (
		mu     sync.Mutex
		seen   []time.Time
		starts []time.Time
	)
	w := Worker{
		Name:     "w",
		Interval: 5 * time.Millisecond,
		Task: func(_ context.Context, since time.Time) error {
			mu.Lock()
			defer mu.Unlock()

			i := len(seen)
			seen = append(seen, since)
			starts = append(starts, time.Now())
			if i == len(results)-1 {
				cancel()
			}

			return results[i]
		},
	}

	done := make(chan error, 1)
	go func() {
		done <- w.Run(ctx, since0)
	}()
	select {
	case err := <-done:
		if err != nil {
			t.Fatal(err)
		}
	case <-time.After(5 * time.Second):
		t.Fatal("worker did not stop")
	}

	mu.Lock()
	defer mu.Unlock()
	if len(seen) != len(results) {
		t.Fatalf("task ran %d times, want %d", len(seen), len(results))
	}
	if !seen[0].Equal(since0) {
		t.Fatalf("first run since = %v, want %v", seen[0], since0)
	}
	// Run 0 succeeded, so runs 1..3 see a checkpoint from run 0's start.
	for i := 1; i <= 3; i++ {
		if !seen[i].Equal(seen[1]) {
			t.Fatalf("run %d since = %v, want unchanged %v", i, seen[i], seen[1])
		}
	}
	if !seen[1].After(since0) || seen[1].After(starts[1]) {
		t.Fatalf(
			"checkpoint after success = %v, want between %v and %v",
			seen[1],
			since0,
			starts[1],
		)
	}
	// Run 3 returned ErrWorkerNoWork, so run 4 sees a newer checkpoint.
	if !seen[4].After(seen[3]) {
		t.Fatalf(
			"checkpoint did not advance after ErrWorkerNoWork: %v -> %v",
			seen[3],
			seen[4],
		)
	}
}

// An in-flight task keeps a live context after the parent is cancelled and is
// allowed to finish.
func TestWorkerRunGracefulShutdown(t *testing.T) {
	ctx, cancel := context.WithCancel(context.Background())
	var workErr error
	w := Worker{
		Name:            "w",
		Interval:        time.Hour,
		ShutdownTimeout: time.Second,
		Task: func(taskCtx context.Context, _ time.Time) error {
			cancel()
			time.Sleep(10 * time.Millisecond)
			workErr = taskCtx.Err()
			return nil
		},
	}

	if err := w.Run(ctx, time.Time{}); err != nil {
		t.Fatal(err)
	}
	if workErr != nil {
		t.Fatalf("task context cancelled before shutdown timeout: %v", workErr)
	}
}

// A task that outlives ShutdownTimeout has its context cancelled.
func TestWorkerRunShutdownTimeout(t *testing.T) {
	ctx, cancel := context.WithCancel(context.Background())
	w := Worker{
		Name:            "w",
		Interval:        time.Hour,
		ShutdownTimeout: 20 * time.Millisecond,
		Task: func(taskCtx context.Context, _ time.Time) error {
			cancel()
			select {
			case <-taskCtx.Done():
				return taskCtx.Err()
			case <-time.After(5 * time.Second):
				return errors.New("task context never cancelled")
			}
		},
	}

	start := time.Now()
	if err := w.Run(ctx, time.Time{}); err != nil {
		t.Fatal(err)
	}
	if d := time.Since(start); d > 2*time.Second {
		t.Fatalf("shutdown took %v", d)
	}
}

func TestWorkerExecute(t *testing.T) {
	since := time.Date(2026, 1, 1, 0, 0, 0, 0, time.UTC)
	logger := slog.Default()
	tests := []struct {
		name    string
		err     error
		advance bool
	}{
		{"success", nil, true},
		{"no work", ErrWorkerNoWork, true},
		{"wrapped no work", fmt.Errorf("%w: nothing", ErrWorkerNoWork), true},
		{"cancelled", context.Canceled, false},
		{"wrapped cancelled", fmt.Errorf("x: %w", context.Canceled), false},
		{"failure", errors.New("boom"), false},
	}
	for _, tc := range tests {
		t.Run(tc.name, func(t *testing.T) {
			w := Worker{
				Task: func(context.Context, time.Time) error {
					return tc.err
				},
			}
			before := time.Now()
			got := w.execute(context.Background(), logger, since)
			if tc.advance {
				if got.Before(before) {
					t.Fatalf(
						"checkpoint %v not advanced to run start (>= %v)", got, before,
					)
				}
			} else if !got.Equal(since) {
				t.Fatalf("checkpoint = %v, want unchanged %v", got, since)
			}
		})
	}
}

func TestAwaitTask(t *testing.T) {
	errFail := errors.New("boom")
	v, err := AwaitTask(context.Background(), func() (int, error) {
		return 42, nil
	})
	if v != 42 || err != nil {
		t.Fatalf("got (%d, %v), want (42, nil)", v, err)
	}
	v, err = AwaitTask(context.Background(), func() (int, error) {
		return 7, errFail
	})
	if v != 7 || !errors.Is(err, errFail) {
		t.Fatalf("got (%d, %v), want (7, %v)", v, err, errFail)
	}

	ctx, cancel := context.WithCancel(context.Background())
	release := make(chan struct{})
	defer close(release)
	go func() {
		time.Sleep(5 * time.Millisecond)
		cancel()
	}()
	v, err = AwaitTask(ctx, func() (int, error) {
		<-release
		return 1, nil
	})
	if v != 0 || !errors.Is(err, context.Canceled) {
		t.Fatalf("got (%d, %v), want (0, context.Canceled)", v, err)
	}
}

func TestAwaitTaskE(t *testing.T) {
	errFail := errors.New("boom")
	if err := AwaitTaskE(context.Background(), func() error {
		return nil
	}); err != nil {
		t.Fatal(err)
	}
	if err := AwaitTaskE(
		context.Background(),
		func() error {
			return errFail
		},
	); !errors.Is(err, errFail) {
		t.Fatalf("got %v, want %v", err, errFail)
	}

	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Millisecond)
	defer cancel()
	release := make(chan struct{})
	defer close(release)
	if err := AwaitTaskE(
		ctx,
		func() error {
			<-release
			return nil
		},
	); !errors.Is(err, context.DeadlineExceeded) {
		t.Fatalf("got %v, want context.DeadlineExceeded", err)
	}
}