package server

import (
	"context"
	"errors"
	"io"
	"strings"
	"testing"
	"time"

	"github.com/stretchr/testify/assert"

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

// scriptJob returns failErr for its first failCalls calls, nil afterwards.
type scriptJob struct {
	failCalls int
	failErr   error
	warnings  []string
	calls     int
}

func (j *scriptJob) Kind() jobKind { return jobMigrateRun }

func (j *scriptJob) Run(context.Context) error {
	j.calls++
	if j.calls <= j.failCalls {
		return j.failErr
	}
	return nil
}

func (j *scriptJob) Warnings() []string { return j.warnings }

func TestRetryJobRun(t *testing.T) {
	partial := exitcode.PartialFailure(errors.New("1 tenant(s) failed: t1"))
	cases := []struct {
		name       string
		failCalls  int
		failErr    error
		retries    int
		wantCalls  int
		wantErrStr string
		wantCode   exitcode.Code
	}{
		{
			name:       "PartialThenSuccess",
			failCalls:  1,
			failErr:    partial,
			retries:    2,
			wantCalls:  2,
			wantErrStr: "",
			wantCode:   exitcode.OK,
		},
		{
			name:       "BudgetExhausted",
			failCalls:  10,
			failErr:    partial,
			retries:    2,
			wantCalls:  3,
			wantErrStr: partial.Error(),
			wantCode:   exitcode.Partial,
		},
		{
			name:       "HardErrorNoRetry",
			failCalls:  10,
			failErr:    errors.New("broken config"),
			retries:    2,
			wantCalls:  1,
			wantErrStr: "broken config",
			wantCode:   exitcode.Generic,
		},
		{
			name:       "CancelledNoRetry",
			failCalls:  10,
			failErr:    exitcode.CancelledFailure(errors.New("stopped")),
			retries:    2,
			wantCalls:  1,
			wantErrStr: "stopped",
			wantCode:   exitcode.Cancelled,
		},
	}
	for _, tc := range cases {
		t.Run(tc.name, func(t *testing.T) {
			inner := &scriptJob{failCalls: tc.failCalls, failErr: tc.failErr}
			job := &retryJob{inner: inner, retries: tc.retries, progress: io.Discard, interval: time.Millisecond}
			err := job.Run(context.Background())
			assert.Equal(t, tc.wantCalls, inner.calls, "inner calls")
			assert.Equal(t, tc.wantCode, exitcode.Classify(err))
			if tc.wantErrStr == "" {
				assert.NoError(t, err)
			} else {
				assert.EqualError(t, err, tc.wantErrStr)
			}
		})
	}
}

func TestRetryJobRetriesZeroRunsOnce(t *testing.T) {
	inner := &scriptJob{failCalls: 10, failErr: exitcode.PartialFailure(errors.New("mismatch"))}
	job := wrapRetry(inner, 0, nil, retryInterval)
	assert.Same(t, inner, job, "wrapRetry with retries=0 must return the bare job")

	assert.Error(t, inner.Run(context.Background()))
	assert.Equal(t, 1, inner.calls)
}

func TestRetryJobCancelDuringSleep(t *testing.T) {
	ctx, cancel := context.WithCancel(context.Background())
	inner := &scriptJob{failCalls: 10, failErr: exitcode.PartialFailure(errors.New("mismatch"))}
	job := &retryJob{inner: inner, retries: 1, progress: io.Discard, interval: time.Minute}

	time.AfterFunc(10*time.Millisecond, cancel)
	started := time.Now()
	err := job.Run(ctx)

	assert.ErrorIs(t, err, context.Canceled)
	assert.Equal(t, 1, inner.calls)
	assert.Less(t, time.Since(started), 30*time.Second, "sleep must be interrupted by cancel")
}

func TestRetryJobMergesWarningsAcrossAttempts(t *testing.T) {
	inner := &scriptJob{
		failCalls: 2,
		failErr:   exitcode.PartialFailure(errors.New("mismatch")),
		warnings:  []string{"tenant=t1 action log unavailable"},
	}
	job := &retryJob{inner: inner, retries: 1, progress: io.Discard, interval: time.Millisecond}

	assert.Error(t, job.Run(context.Background()))
	assert.Equal(t, 2, inner.calls)
	assert.Equal(t, []string{"tenant=t1 action log unavailable"}, job.Warnings())
}

func TestRetryJobAnnouncesRetry(t *testing.T) {
	var progress strings.Builder
	inner := &scriptJob{failCalls: 1, failErr: exitcode.PartialFailure(errors.New("mismatch"))}
	job := &retryJob{inner: inner, retries: 1, progress: &progress, interval: time.Millisecond}

	assert.NoError(t, job.Run(context.Background()))
	assert.Contains(t, progress.String(), "retry 1/1 in 1ms")
}