package ssh

import (
	"context"
	"errors"
	"testing"
	"time"
)

func TestRunWithRetryStopsAfterConfiguredRetries(t *testing.T) {
	wantErr := errors.New("ssh exited")
	calls := 0
	err := runWithRetry(context.Background(), 2, 0, func() error {
		calls++
		return wantErr
	}, nil)

	if !errors.Is(err, wantErr) {
		t.Fatalf("runWithRetry() error = %v, want %v", err, wantErr)
	}
	if calls != 3 {
		t.Errorf("run() calls = %d, want 3", calls)
	}
}

func TestRunWithRetryDoesNotRetryWhenDisabled(t *testing.T) {
	wantErr := errors.New("ssh exited")
	calls := 0
	err := runWithRetry(context.Background(), 0, 0, func() error {
		calls++
		return wantErr
	}, nil)

	if !errors.Is(err, wantErr) {
		t.Fatalf("runWithRetry() error = %v, want %v", err, wantErr)
	}
	if calls != 1 {
		t.Errorf("run() calls = %d, want 1", calls)
	}
}

func TestRunWithRetryRetriesIndefinitelyUntilSuccess(t *testing.T) {
	calls := 0
	err := runWithRetry(context.Background(), -1, 0, func() error {
		calls++
		if calls < 4 {
			return errors.New("temporary ssh failure")
		}
		return nil
	}, nil)

	if err != nil {
		t.Fatalf("runWithRetry() error = %v, want nil", err)
	}
	if calls != 4 {
		t.Errorf("run() calls = %d, want 4", calls)
	}
}

func TestRunWithRetryStopsWaitingWhenContextCancelled(t *testing.T) {
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	firstCall := make(chan struct{})
	done := make(chan error, 1)
	go func() {
		done <- runWithRetry(ctx, -1, time.Hour, func() error {
			close(firstCall)
			return errors.New("temporary ssh failure")
		}, nil)
	}()

	<-firstCall
	cancel()

	select {
	case err := <-done:
		if !errors.Is(err, context.Canceled) {
			t.Errorf("runWithRetry() error = %v, want context cancellation", err)
		}
	case <-time.After(time.Second):
		t.Fatal("runWithRetry() did not stop after context cancellation")
	}
}