package dispatcher

import (
	"context"
	"errors"
	"testing"

	"github.com/pingcap/failpoint"
	"github.com/stretchr/testify/require"

	"github.com/tidbcloud/AgentTeam/desktop/internal/agent"
	"github.com/tidbcloud/AgentTeam/desktop/internal/agent/drivercontract"
	"github.com/tidbcloud/AgentTeam/desktop/internal/agent/runtimeconfig"
	"github.com/tidbcloud/AgentTeam/desktop/internal/types"
)

func TestCodexLaunchPersistsUnsafeEpochBeforeNewProcessSpawn(t *testing.T) {
	stubRuntimeBinaries(t, "codex")
	require.NoError(t, failpoint.Enable("dispatcher-new-agent-process", "return(true)"))
	t.Cleanup(func() {
		SetNewAgentProcessHook(nil)
		require.NoError(t, failpoint.Disable("dispatcher-new-agent-process"))
	})

	for _, tc := range []struct {
		name      string
		sessionID string
		prompt    string
	}{
		{name: "fresh initial turn", prompt: "initial fresh body"},
		{name: "resume initial turn", sessionID: "thread-old", prompt: "initial resumed body"},
		{name: "fresh no-prompt launch"},
	} {
		t.Run(tc.name, func(t *testing.T) {
			workDir := t.TempDir()
			require.NoError(t, saveSessionResumeState(workDir, sessionResumeState{
				LastSessionID:              "thread-old",
				SessionResumeSafe:          true,
				SessionResumeSafetyEpoch:   41,
				SessionResumeSafetyVersion: 9,
				LastStartConfig: &types.AgentStartConfig{
					Name:      "Codex Agent",
					Runtime:   string(runtimeconfig.TypeCodex),
					SessionID: "thread-old",
				},
			}))

			spawnObserved := false
			SetNewAgentProcessHook(func(
				_ context.Context,
				_ drivercontract.Driver,
				cfg drivercontract.LaunchConfig,
				options agent.ProcessOptions,
			) (*agent.Process, error) {
				spawnObserved = true
				state, err := loadSessionResumeState(workDir)
				require.NoError(t, err)
				require.False(t, state.SessionResumeSafe, "Spawn must never observe the previous safe candidate")
				require.Equal(t, uint64(42), state.SessionResumeSafetyEpoch)
				require.Zero(t, state.SessionResumeSafetyVersion)
				require.Equal(t, tc.sessionID, state.LastSessionID, "fresh launch must clear the old resume candidate")
				require.NotNil(t, state.LastStartConfig)
				require.Equal(t, tc.sessionID, state.LastStartConfig.SessionID)
				require.Equal(t, tc.sessionID, cfg.SessionID)
				require.Equal(t, tc.prompt, testLaunchSurface(cfg, "prompt"))
				return nil, errors.New("injected spawn failure after observing durable baseline")
			})

			d := mustNewDispatcher("http://server.test")
			t.Cleanup(d.Close)
			sess := &interactiveSession{
				agentID: "agent-codex-pre-spawn",
				agent: agent.Config{
					ID:   "agent-codex-pre-spawn",
					Name: "Codex Agent",
				},
				runtime: runtimeconfig.Runtime{Type: runtimeconfig.TypeCodex, Binary: "codex"},
				binding: conversationBinding{channelID: "ch-1"},
				workDir: workDir,
			}
			_, err := d.launch(context.Background(), &processLaunchContext{
				sess:     sess,
				agent:    sess.agent,
				runtime:  sess.runtime,
				driver:   &stubDriver{},
				binding:  sess.binding,
				startCfg: testLaunchConfig(drivercontract.LaunchConfig{AgentID: sess.agentID, RuntimeType: runtimeconfig.TypeCodex, SessionID: tc.sessionID, WorkDir: workDir, Binary: "codex", PermissionMode: types.AgentPermissionModeFullAccess}, testLaunchSurfaces{Prompt: tc.prompt, FreshPrompt: tc.prompt, ResumePrompt: tc.prompt}),
			})

			require.ErrorContains(t, err, "injected spawn failure")
			require.True(t, spawnObserved)
			state, loadErr := loadSessionResumeState(workDir)
			require.NoError(t, loadErr)
			require.False(t, state.SessionResumeSafe, "spawn failure leaves the new epoch fail-closed")
			require.Equal(t, uint64(42), state.SessionResumeSafetyEpoch)
		})
	}
}