package orchestrator

import (
	"context"
	"sync"
	"testing"
	"time"

	"github.com/stretchr/testify/require"

	"github.com/tidbcloud/AgentTeam/internal/store/workflow/run"
)

// observationReaderFake implements RuntimeObservationReader with a deterministic
// barrier: the first acquire returns the pre-barrier observation and an acquired
// release; every later acquire returns the post-barrier observation. It records
// whether the release was called, proving the runner holds the guard across the
// tick and releases it on every path.
type observationReaderFake struct {
	mu              sync.Mutex
	preWorking      bool
	preKnown        bool
	postWorking     bool
	postKnown       bool
	acquires        int
	released        bool
	releasedAfter   int // acquires count at release time
	acquireReleases []int
}

func (reader *observationReaderFake) AcquireRuntimeObservation(_ context.Context, _ string, _ string) (run.RecoveryObservation, func(), bool) {
	reader.mu.Lock()
	defer reader.mu.Unlock()
	reader.acquires++
	if reader.acquires <= 1 {
		return run.RecoveryObservation{
				SessionGeneration: "gen_1", Connected: true, StateKnown: true, State: "idle",
				Revision: 1, HighWater: 1, Working: reader.preWorking, Known: reader.preKnown,
			}, func() {
				reader.mu.Lock()
				reader.released = true
				reader.releasedAfter = reader.acquires
				reader.mu.Unlock()
			}, true
	}
	return run.RecoveryObservation{
			SessionGeneration: "gen_2", Connected: true, StateKnown: true, State: "active",
			Revision: 2, HighWater: 2, Working: reader.postWorking, Known: reader.postKnown,
		}, func() {
			reader.mu.Lock()
			reader.released = true
			reader.releasedAfter = reader.acquires
			reader.mu.Unlock()
		}, true
}

func (reader *observationReaderFake) RuntimeWorking(_ context.Context, _ string, _ string) (bool, bool) {
	reader.mu.Lock()
	defer reader.mu.Unlock()
	working := reader.preWorking
	if reader.acquires > 1 {
		working = reader.postWorking
	}
	return working, reader.preKnown || reader.postKnown
}

// TestRecoveryRunnerPassesImmutableObservationAndReleasesGuard proves the
// RecoveryRunner acquires the observation through the production observation
// reader, passes the immutable observation (with the working/known judgment) to
// RecoveryTickParams, and releases the guard after the tick returns.
func TestRecoveryRunnerPassesImmutableObservationAndReleasesGuard(t *testing.T) {
	effect := recoveryEffect(t, "wri_fence", "wfx_fence")
	reader := &observationReaderFake{preWorking: false, preKnown: true, postWorking: true, postKnown: true}
	var captured run.RecoveryTickParams
	repository := &fakeRecoveryRepository{
		candidates: []run.Effect{effect},
		tick: func(params run.RecoveryTickParams) (run.RecoveryTickResult, error) {
			captured = params
			return run.RecoveryTickResult{Effect: effect, Send: false}, nil
		},
	}
	writer := &recordingRecoveryWriter{}
	runner, err := NewRecoveryRunner(repository, writer, RecoveryRunnerOptions{
		WorkerID:       "recovery-fence",
		Now:            func() time.Time { return time.Date(2026, 8, 17, 12, 0, 0, 0, time.UTC) },
		RuntimeWorking: reader,
	})
	require.NoError(t, err)

	outcome, err := runner.RecoverOnce(context.Background())
	require.NoError(t, err)
	require.Equal(t, 1, outcome.Candidates)
	require.Zero(t, outcome.Working, "an idle observation is not counted working")
	require.Zero(t, outcome.Sent)

	// The immutable observation must be carried into the tick params.
	require.True(t, captured.Observation.Known, "the observation must carry the known judgment")
	require.False(t, captured.Observation.Working, "the observation must carry the idle judgment")
	require.Equal(t, "gen_1", captured.Observation.SessionGeneration, "the observation must carry the fenced session generation")
	require.Equal(t, int64(1), captured.Observation.Revision, "the observation must carry the fenced revision")

	// The guard must be released exactly once after the tick.
	reader.mu.Lock()
	released := reader.released
	releasedAfter := reader.releasedAfter
	reader.mu.Unlock()
	require.True(t, released, "the machine observation guard must be released after the tick")
	require.Equal(t, 1, releasedAfter, "the guard must be released after the single acquire")
}

// TestRecoveryRunnerSkipsUnknownObservationAndReleasesGuard proves an unknown
// observation (known=false) causes the runner to skip the effect with zero
// transition/write AND release the guard on the skip path.
func TestRecoveryRunnerSkipsUnknownObservationAndReleasesGuard(t *testing.T) {
	effect := recoveryEffect(t, "wri_fence_unknown", "wfx_fence_unknown")
	reader := &observationReaderFake{preWorking: false, preKnown: false, postWorking: false, postKnown: false}
	called := false
	repository := &fakeRecoveryRepository{
		candidates: []run.Effect{effect},
		tick: func(params run.RecoveryTickParams) (run.RecoveryTickResult, error) {
			called = true
			return run.RecoveryTickResult{Effect: effect, Send: false}, nil
		},
	}
	writer := &recordingRecoveryWriter{}
	runner, err := NewRecoveryRunner(repository, writer, RecoveryRunnerOptions{
		WorkerID:       "recovery-fence-unknown",
		Now:            func() time.Time { return time.Date(2026, 8, 17, 12, 0, 0, 0, time.UTC) },
		RuntimeWorking: reader,
	})
	require.NoError(t, err)

	outcome, err := runner.RecoverOnce(context.Background())
	require.NoError(t, err)
	require.Equal(t, 1, outcome.Unknown, "unknown observation is skipped")
	require.False(t, called, "the tick must not be reached for an unknown observation")

	reader.mu.Lock()
	released := reader.released
	reader.mu.Unlock()
	require.True(t, released, "the guard must be released on the unknown skip path")
}