package execution

import (
	"context"
	"errors"
	"sync/atomic"
	"testing"
	"time"

	"github.com/stretchr/testify/require"

	executionstore "github.com/tidbcloud/AgentTeam/internal/store/role/execution"
)

func TestNewWorkerRejectsMissingDependencies(t *testing.T) {
	valid := workerTestOptions()
	for _, testCase := range []struct {
		name   string
		mutate func(*WorkerOptions)
	}{
		{
			name: "workspace source",
			mutate: func(options *WorkerOptions) {
				options.WorkspaceIDs = nil
			},
		},
		{
			name: "claimer",
			mutate: func(options *WorkerOptions) {
				options.ClaimRunnable = nil
			},
		},
		{
			name: "executor",
			mutate: func(options *WorkerOptions) {
				options.Execute = nil
			},
		},
	} {
		t.Run(testCase.name, func(t *testing.T) {
			options := valid
			testCase.mutate(&options)

			worker, err := NewWorker(options)

			require.Nil(t, worker)
			require.EqualError(t, err, "role management worker dependencies are incomplete")
		})
	}
}

func TestNewWorkerNormalizesDefaults(t *testing.T) {
	for _, testCase := range []struct {
		name   string
		mutate func(*WorkerOptions)
	}{
		{
			name: "zero values",
			mutate: func(options *WorkerOptions) {
				options.Now = nil
				options.Interval = 0
				options.LeaseTTL = 0
				options.Limit = 0
			},
		},
		{
			name: "negative values",
			mutate: func(options *WorkerOptions) {
				options.Now = nil
				options.Interval = -time.Second
				options.LeaseTTL = -time.Second
				options.Limit = -1
			},
		},
	} {
		t.Run(testCase.name, func(t *testing.T) {
			options := workerTestOptions()
			testCase.mutate(&options)

			worker, err := NewWorker(options)

			require.NoError(t, err)
			require.NotNil(t, worker.now)
			require.False(t, worker.now().IsZero())
			require.Equal(t, 5*time.Second, worker.interval)
			require.Equal(t, 2*time.Minute, worker.leaseTTL)
			require.Equal(t, 1, worker.limit)
		})
	}
}

func TestWorkerRunProcessesBeforeFirstTick(t *testing.T) {
	called := make(chan struct{}, 1)
	options := workerTestOptions()
	options.Interval = time.Hour
	options.WorkspaceIDs = func(context.Context) ([]string, error) {
		called <- struct{}{}
		return nil, nil
	}
	worker, err := NewWorker(options)
	require.NoError(t, err)
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	go worker.Run(ctx)

	requireSignal(t, called, "worker pass before the one-hour first tick")
	cancel()
}

func TestWorkerRunProcessesLaterTickAfterImmediatePass(t *testing.T) {
	var sourceCalls atomic.Int32
	immediatePass := make(chan struct{}, 1)
	releaseImmediate := make(chan struct{})
	secondPass := make(chan struct{}, 1)
	options := workerTestOptions()
	options.Interval = 10 * time.Millisecond
	options.WorkspaceIDs = func(context.Context) ([]string, error) {
		switch sourceCalls.Add(1) {
		case 1:
			immediatePass <- struct{}{}
			<-releaseImmediate
		case 2:
			secondPass <- struct{}{}
		}
		return nil, nil
	}
	worker, err := NewWorker(options)
	require.NoError(t, err)
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	go worker.Run(ctx)

	requireSignal(t, immediatePass, "immediate worker pass")
	require.Equal(t, int32(1), sourceCalls.Load(),
		"later tick must be observed only after synchronizing the immediate pass")
	close(releaseImmediate)
	requireSignal(t, secondPass, "later worker tick")
	cancel()
}

func TestWorkerRunCancelledBeforeStartDoesNoWork(t *testing.T) {
	var sourceCalls atomic.Int32
	options := workerTestOptions()
	options.WorkspaceIDs = func(context.Context) ([]string, error) {
		sourceCalls.Add(1)
		return []string{"ws_a"}, nil
	}
	worker, err := NewWorker(options)
	require.NoError(t, err)
	ctx, cancel := context.WithCancel(context.Background())
	cancel()

	worker.Run(ctx)

	require.Zero(t, sourceCalls.Load())
}

func TestWorkerProcessOnceRefreshesAndNormalizesWorkspaceIDs(t *testing.T) {
	now := time.Date(2026, 8, 5, 3, 40, 0, 123456789, time.FixedZone("CST", 8*60*60))
	var sourceCalls int
	var claimParams []executionstore.ClaimParams
	options := workerTestOptions()
	options.Now = func() time.Time { return now }
	options.WorkspaceIDs = func(context.Context) ([]string, error) {
		sourceCalls++
		return []string{" ws_b ", "", "ws_a", "ws_b", "  ", "ws_c"}, nil
	}
	options.ClaimRunnable = func(
		_ context.Context,
		params executionstore.ClaimParams,
	) ([]executionstore.ClaimedExecution, error) {
		claimParams = append(claimParams, params)
		return nil, nil
	}
	worker, err := NewWorker(options)
	require.NoError(t, err)

	worker.ProcessOnce(context.Background())
	worker.ProcessOnce(context.Background())

	require.Equal(t, 2, sourceCalls, "every pass must refresh owned workspaces")
	require.Len(t, claimParams, 6)
	for index, workspaceID := range []string{
		"ws_a", "ws_b", "ws_c", "ws_a", "ws_b", "ws_c",
	} {
		require.Equal(t, workspaceID, claimParams[index].WorkspaceID)
		require.Equal(t, now.UTC(), claimParams[index].Now)
		require.Equal(t, 2*time.Minute, claimParams[index].LeaseTTL)
		require.Equal(t, 1, claimParams[index].Limit)
	}
}

func TestWorkerSourceErrorSkipsClaimsAndLaterPassRecovers(t *testing.T) {
	sourceErr := errors.New("workspace source unavailable")
	var sourceCalls, claimCalls int
	options := workerTestOptions()
	options.WorkspaceIDs = func(context.Context) ([]string, error) {
		sourceCalls++
		if sourceCalls == 1 {
			return nil, sourceErr
		}
		return []string{"ws_recovered"}, nil
	}
	options.ClaimRunnable = func(
		_ context.Context,
		params executionstore.ClaimParams,
	) ([]executionstore.ClaimedExecution, error) {
		claimCalls++
		require.Equal(t, "ws_recovered", params.WorkspaceID)
		return nil, nil
	}
	worker, err := NewWorker(options)
	require.NoError(t, err)

	worker.ProcessOnce(context.Background())
	require.Zero(t, claimCalls)
	worker.ProcessOnce(context.Background())

	require.Equal(t, 2, sourceCalls)
	require.Equal(t, 1, claimCalls)
}

func TestWorkerEmptyClaimDoesNotExecute(t *testing.T) {
	var executeCalls int
	options := workerTestOptions()
	options.WorkspaceIDs = func(context.Context) ([]string, error) {
		return []string{"ws_empty"}, nil
	}
	options.Execute = func(
		context.Context,
		executionstore.ClaimedExecution,
		time.Time,
	) error {
		executeCalls++
		return nil
	}
	worker, err := NewWorker(options)
	require.NoError(t, err)

	worker.ProcessOnce(context.Background())

	require.Zero(t, executeCalls)
}

func TestWorkerExecutesEachClaimOnceWithWorkspaceClaimNow(t *testing.T) {
	times := []time.Time{
		time.Date(2026, 8, 5, 3, 50, 0, 0, time.UTC),
		time.Date(2026, 8, 5, 3, 51, 0, 0, time.UTC),
	}
	var nowCalls int
	var calls []workerExecuteCall
	options := workerTestOptions()
	options.Now = func() time.Time {
		value := times[nowCalls]
		nowCalls++
		return value
	}
	options.WorkspaceIDs = func(context.Context) ([]string, error) {
		return []string{"ws_b", "ws_a"}, nil
	}
	options.ClaimRunnable = func(
		_ context.Context,
		params executionstore.ClaimParams,
	) ([]executionstore.ClaimedExecution, error) {
		return []executionstore.ClaimedExecution{
			workerClaim(params.WorkspaceID, "approval_1"),
			workerClaim(params.WorkspaceID, "approval_2"),
		}, nil
	}
	options.Execute = func(
		_ context.Context,
		claim executionstore.ClaimedExecution,
		claimNow time.Time,
	) error {
		calls = append(calls, workerExecuteCall{claim: claim, now: claimNow})
		return nil
	}
	worker, err := NewWorker(options)
	require.NoError(t, err)

	worker.ProcessOnce(context.Background())

	require.Equal(t, 2, nowCalls, "Now is captured once per workspace")
	require.Len(t, calls, 4)
	for index, call := range calls {
		wantWorkspace := "ws_a"
		wantNow := times[0]
		if index >= 2 {
			wantWorkspace = "ws_b"
			wantNow = times[1]
		}
		require.Equal(t, wantWorkspace, call.claim.Execution.WorkspaceID)
		require.Equal(t, wantNow, call.now)
	}
}

func TestWorkerClaimAndExecuteErrorsDoNotStarveLaterWorkspaces(t *testing.T) {
	claimErr := errors.New("claim failed")
	executeErr := errors.New("execute failed")
	var executed []string
	options := workerTestOptions()
	options.WorkspaceIDs = func(context.Context) ([]string, error) {
		return []string{"ws_claim_error", "ws_execute_error", "ws_later"}, nil
	}
	options.ClaimRunnable = func(
		_ context.Context,
		params executionstore.ClaimParams,
	) ([]executionstore.ClaimedExecution, error) {
		if params.WorkspaceID == "ws_claim_error" {
			return nil, claimErr
		}
		return []executionstore.ClaimedExecution{
			workerClaim(params.WorkspaceID, "approval"),
		}, nil
	}
	options.Execute = func(
		_ context.Context,
		claim executionstore.ClaimedExecution,
		_ time.Time,
	) error {
		workspaceID := claim.Execution.WorkspaceID
		executed = append(executed, workspaceID)
		if workspaceID == "ws_execute_error" {
			return executeErr
		}
		return nil
	}
	worker, err := NewWorker(options)
	require.NoError(t, err)

	worker.ProcessOnce(context.Background())

	require.Equal(t, []string{"ws_execute_error", "ws_later"}, executed)
}

func TestWorkerCancellationStopsLaterDispatch(t *testing.T) {
	for _, testCase := range []struct {
		name   string
		cancel func(*workerCancelProbe, context.CancelFunc)
	}{
		{
			name: "during claim",
			cancel: func(probe *workerCancelProbe, cancel context.CancelFunc) {
				probe.cancelClaim = cancel
			},
		},
		{
			name: "during execute",
			cancel: func(probe *workerCancelProbe, cancel context.CancelFunc) {
				probe.cancelExecute = cancel
			},
		},
	} {
		t.Run(testCase.name, func(t *testing.T) {
			probe := &workerCancelProbe{}
			ctx, cancel := context.WithCancel(context.Background())
			testCase.cancel(probe, cancel)
			worker, err := NewWorker(probe.options())
			require.NoError(t, err)

			worker.ProcessOnce(ctx)

			require.Equal(t, 1, probe.claimCalls)
			if testCase.name == "during claim" {
				require.Zero(t, probe.executeCalls)
			} else {
				require.Equal(t, 1, probe.executeCalls)
			}
		})
	}
}

type workerExecuteCall struct {
	claim executionstore.ClaimedExecution
	now   time.Time
}

type workerCancelProbe struct {
	claimCalls    int
	executeCalls  int
	cancelClaim   context.CancelFunc
	cancelExecute context.CancelFunc
}

func (probe *workerCancelProbe) options() WorkerOptions {
	options := workerTestOptions()
	options.WorkspaceIDs = func(context.Context) ([]string, error) {
		return []string{"ws_first", "ws_later"}, nil
	}
	options.ClaimRunnable = func(
		_ context.Context,
		params executionstore.ClaimParams,
	) ([]executionstore.ClaimedExecution, error) {
		probe.claimCalls++
		if probe.cancelClaim != nil {
			probe.cancelClaim()
		}
		return []executionstore.ClaimedExecution{
			workerClaim(params.WorkspaceID, "approval"),
		}, nil
	}
	options.Execute = func(
		_ context.Context,
		_ executionstore.ClaimedExecution,
		_ time.Time,
	) error {
		probe.executeCalls++
		if probe.cancelExecute != nil {
			probe.cancelExecute()
		}
		return nil
	}
	return options
}

func workerTestOptions() WorkerOptions {
	return WorkerOptions{
		WorkspaceIDs: func(context.Context) ([]string, error) {
			return nil, nil
		},
		ClaimRunnable: func(
			context.Context,
			executionstore.ClaimParams,
		) ([]executionstore.ClaimedExecution, error) {
			return nil, nil
		},
		Execute: func(
			context.Context,
			executionstore.ClaimedExecution,
			time.Time,
		) error {
			return nil
		},
		Now:      time.Now,
		Interval: time.Hour,
	}
}

func workerClaim(workspaceID, approvalID string) executionstore.ClaimedExecution {
	return executionstore.ClaimedExecution{
		Execution: executionstore.Execution{
			WorkspaceID: workspaceID,
			ApprovalID:  approvalID,
			Status:      executionstore.ExecutionRunning,
		},
		Fence: executionstore.ExecutionFence{
			WorkspaceID: workspaceID,
			ApprovalID:  approvalID,
			LeaseToken:  "lease_" + approvalID,
		},
	}
}

func requireSignal(t testing.TB, signal <-chan struct{}, description string) {
	t.Helper()
	select {
	case <-signal:
	case <-time.After(time.Second):
		t.Fatalf("timed out waiting for %s", description)
	}
}