package agent

import (
	"context"
	"testing"

	"github.com/stretchr/testify/require"

	"github.com/tidbcloud/AgentTeam/internal/core/rewindcaps"
)

func TestCodexRestoreRecoveryPointResumesPersistedThreadBeforeRollbackAfterRestart(t *testing.T) {
	driver, writer := newTestCodexDriver(t, func(req capturedRPCRequest) appServerEnvelope {
		switch req.Method {
		case "thread/resume":
			require.Equal(t, "thread-1", req.Params["threadId"])
			return appServerEnvelope{Result: rawRPCResult(t, appServerThreadResult{
				Thread: appServerThread{ID: "thread-1", Status: appServerThreadStatus{Type: "idle"}},
			})}
		case "thread/read":
			require.Equal(t, "thread-1", req.Params["threadId"])
			require.Equal(t, true, req.Params["includeTurns"])
			return appServerEnvelope{Result: rawRPCResult(t, appServerThreadResult{
				Thread: appServerThread{
					ID:     "thread-1",
					Status: appServerThreadStatus{Type: "idle"},
					Turns: []appServerTurn{
						{ID: "turn-1", Status: appServerTurnStatus{Type: "completed"}},
						{ID: "turn-2", Status: appServerTurnStatus{Type: "completed"}},
					},
				},
			})}
		case "thread/rollback":
			require.Equal(t, "thread-1", req.Params["threadId"])
			require.Equal(t, float64(1), req.Params["numTurns"])
			return appServerEnvelope{Result: rawRPCResult(t, appServerThreadResult{
				Thread: appServerThread{
					ID:     "thread-1",
					Status: appServerThreadStatus{Type: "idle"},
					Turns: []appServerTurn{
						{ID: "turn-1", Status: appServerTurnStatus{Type: "completed"}},
					},
				},
			})}
		default:
			t.Fatalf("unexpected request %s", req.Method)
			return appServerEnvelope{}
		}
	})
	recoveryRef, err := rewindcaps.EncodeCodexAppServerRecoveryRef("thread-1", 1)
	require.NoError(t, err)

	err = driver.RestoreRecoveryPoint(context.Background(), rewindcaps.StrategyCodexAppServerTurns, recoveryRef)

	require.NoError(t, err)
	require.Equal(t, []string{"thread/resume", "thread/read", "thread/rollback"}, capturedRPCMethods(writer.requests))
}