package agent

import (
	"context"
	"testing"
	"time"

	"github.com/stretchr/testify/require"

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

const codexRecoveryInterruptTurnID = "turn-interrupted"

func codexRecoveryHistory(count int) []appServerTurn {
	turns := make([]appServerTurn, 0, count)
	for i := 0; i < count; i++ {
		turns = append(turns, appServerTurn{
			ID:     "turn-history-" + string(rune('1'+i)),
			Status: appServerTurnStatus{Type: "completed"},
		})
	}
	return turns
}

func codexRecoveryThread(turns []appServerTurn) appServerThread {
	return appServerThread{
		ID:     "thread-1",
		Status: appServerThreadStatus{Type: "idle"},
		Turns:  turns,
	}
}

func codexRecoveryMaterializedInterruptThread() appServerThread {
	turns := codexRecoveryHistory(2)
	turns = append(turns, appServerTurn{
		ID:     codexRecoveryInterruptTurnID,
		Status: appServerTurnStatus{Type: "interrupted"},
	})
	return codexRecoveryThread(turns)
}

func codexRecoveryRefAtTwo(t *testing.T) string {
	t.Helper()
	recoveryRef, err := rewindcaps.EncodeCodexAppServerRecoveryRef("thread-1", 2)
	require.NoError(t, err)
	return recoveryRef
}

func seedCodexRecoveryActiveTurn(driver *CodexAppServerDriver) {
	driver.threadID = "thread-1"
	driver.threadStatus = appServerThreadStatus{Type: "active"}
	driver.activeTurnID = codexRecoveryInterruptTurnID
	driver.completedTurnCount = 2
}

func TestCodexRestoreWaitsForDelayedInterruptedTurnBeforeRollback(t *testing.T) {
	for _, tc := range []struct {
		name             string
		initialTurnCount int
	}{
		{name: "current below target", initialTurnCount: 1},
		{name: "current exactly target", initialTurnCount: 2},
	} {
		t.Run(tc.name, func(t *testing.T) {
			readCount := 0
			rollbackCount := 0
			driver, writer := newTestCodexDriver(t, func(req capturedRPCRequest) appServerEnvelope {
				switch req.Method {
				case "turn/interrupt":
					require.Equal(t, codexRecoveryInterruptTurnID, req.Params["turnId"])
					return appServerEnvelope{Result: rawRPCResult(t, map[string]any{})}
				case "thread/read":
					readCount++
					var thread appServerThread
					switch readCount {
					case 1:
						// InterruptActiveRun has observed idle, but Codex has not yet
						// appended the terminal interrupted turn.
						thread = codexRecoveryThread(codexRecoveryHistory(2))
					case 2:
						// Restore's coarse exclude-turns read.
						thread = codexRecoveryThread(nil)
					case 3:
						// The first authoritative includeTurns snapshot is either below
						// or exactly at the requested checkpoint; neither is a no-op.
						thread = codexRecoveryThread(codexRecoveryHistory(tc.initialTurnCount))
					default:
						thread = codexRecoveryMaterializedInterruptThread()
					}
					return appServerEnvelope{Result: rawRPCResult(t, appServerThreadResult{Thread: thread})}
				case "thread/rollback":
					rollbackCount++
					require.Equal(t, float64(1), req.Params["numTurns"])
					return appServerEnvelope{Result: rawRPCResult(t, appServerThreadResult{
						Thread: codexRecoveryThread(codexRecoveryHistory(2)),
					})}
				default:
					t.Fatalf("unexpected method %s", req.Method)
					return appServerEnvelope{}
				}
			})
			seedCodexRecoveryActiveTurn(driver)

			require.NoError(t, driver.InterruptActiveRun(context.Background(), "thread-1"))
			err := driver.RestoreRecoveryPoint(
				context.Background(), rewindcaps.StrategyCodexAppServerTurns, codexRecoveryRefAtTwo(t),
			)

			require.NoError(t, err)
			require.Equal(t, 1, rollbackCount)
			require.GreaterOrEqual(t, readCount, 4, "restore must poll beyond the stale authoritative snapshot")
			require.Equal(t, "thread/rollback", writer.requests[len(writer.requests)-1].Method)
			driver.mu.Lock()
			snapshot := driver.lastActiveInterrupt
			driver.mu.Unlock()
			require.Empty(t, snapshot.threadID)
		})
	}
}

func TestCodexRestoreConsumesAlreadyVisibleInterruptedTurnWithoutPolling(t *testing.T) {
	readCount := 0
	rollbackCount := 0
	driver, _ := newTestCodexDriver(t, func(req capturedRPCRequest) appServerEnvelope {
		switch req.Method {
		case "turn/interrupt":
			return appServerEnvelope{Result: rawRPCResult(t, map[string]any{})}
		case "thread/read":
			readCount++
			thread := codexRecoveryMaterializedInterruptThread()
			if req.Params["includeTurns"] == false {
				thread = codexRecoveryThread(nil)
			}
			return appServerEnvelope{Result: rawRPCResult(t, appServerThreadResult{Thread: thread})}
		case "thread/rollback":
			rollbackCount++
			return appServerEnvelope{Result: rawRPCResult(t, appServerThreadResult{
				Thread: codexRecoveryThread(codexRecoveryHistory(2)),
			})}
		default:
			t.Fatalf("unexpected method %s", req.Method)
			return appServerEnvelope{}
		}
	})
	seedCodexRecoveryActiveTurn(driver)
	require.NoError(t, driver.InterruptActiveRun(context.Background(), "thread-1"))

	err := driver.RestoreRecoveryPoint(
		context.Background(), rewindcaps.StrategyCodexAppServerTurns, codexRecoveryRefAtTwo(t),
	)

	require.NoError(t, err)
	require.Equal(t, 1, rollbackCount)
	require.Equal(t, 3, readCount, "materialized initial history must not add a polling read")
}

func TestCodexRestoreInterruptedTurnWaitCancellationRetainsSnapshotForRetry(t *testing.T) {
	readCount := 0
	rollbackCount := 0
	driver, _ := newTestCodexDriver(t, func(req capturedRPCRequest) appServerEnvelope {
		switch req.Method {
		case "turn/interrupt":
			return appServerEnvelope{Result: rawRPCResult(t, map[string]any{})}
		case "thread/read":
			readCount++
			thread := codexRecoveryThread(codexRecoveryHistory(2))
			if readCount == 2 || req.Params["includeTurns"] == false {
				thread = codexRecoveryThread(nil)
			}
			if readCount >= 4 && req.Params["includeTurns"] == true {
				thread = codexRecoveryMaterializedInterruptThread()
			}
			return appServerEnvelope{Result: rawRPCResult(t, appServerThreadResult{Thread: thread})}
		case "thread/rollback":
			rollbackCount++
			return appServerEnvelope{Result: rawRPCResult(t, appServerThreadResult{
				Thread: codexRecoveryThread(codexRecoveryHistory(2)),
			})}
		default:
			t.Fatalf("unexpected method %s", req.Method)
			return appServerEnvelope{}
		}
	})
	seedCodexRecoveryActiveTurn(driver)
	require.NoError(t, driver.InterruptActiveRun(context.Background(), "thread-1"))

	canceledCtx, cancel := context.WithTimeout(context.Background(), 10*time.Millisecond)
	defer cancel()
	err := driver.RestoreRecoveryPoint(
		canceledCtx, rewindcaps.StrategyCodexAppServerTurns, codexRecoveryRefAtTwo(t),
	)
	require.ErrorIs(t, err, context.DeadlineExceeded)
	driver.mu.Lock()
	retained := driver.lastActiveInterrupt
	driver.mu.Unlock()
	require.Equal(t, codexRecoveryInterruptTurnID, retained.interruptedTurnID)
	require.Zero(t, rollbackCount)

	err = driver.RestoreRecoveryPoint(
		context.Background(), rewindcaps.StrategyCodexAppServerTurns, codexRecoveryRefAtTwo(t),
	)
	require.NoError(t, err)
	require.Equal(t, 1, rollbackCount)
	driver.mu.Lock()
	cleared := driver.lastActiveInterrupt
	driver.mu.Unlock()
	require.Empty(t, cleared.threadID)
}

func TestCodexRestoreInterruptedTurnVisibilityExpiryFailsClosedAndClearsSnapshot(t *testing.T) {
	readCount := 0
	rollbackCount := 0
	driver, _ := newTestCodexDriver(t, func(req capturedRPCRequest) appServerEnvelope {
		switch req.Method {
		case "turn/interrupt":
			return appServerEnvelope{Result: rawRPCResult(t, map[string]any{})}
		case "thread/read":
			readCount++
			return appServerEnvelope{Result: rawRPCResult(t, appServerThreadResult{
				Thread: codexRecoveryThread(codexRecoveryHistory(2)),
			})}
		case "thread/rollback":
			rollbackCount++
			return appServerEnvelope{Result: rawRPCResult(t, appServerThreadResult{})}
		default:
			t.Fatalf("unexpected method %s", req.Method)
			return appServerEnvelope{}
		}
	})
	seedCodexRecoveryActiveTurn(driver)
	require.NoError(t, driver.InterruptActiveRun(context.Background(), "thread-1"))
	driver.mu.Lock()
	driver.lastActiveInterrupt.recordedAt = time.Now().Add(-codexActiveInterruptSnapshotTTL - time.Millisecond)
	driver.mu.Unlock()

	err := driver.RestoreRecoveryPoint(
		context.Background(), rewindcaps.StrategyCodexAppServerTurns, codexRecoveryRefAtTwo(t),
	)

	require.Error(t, err)
	require.Contains(t, err.Error(), "expired")
	require.Zero(t, rollbackCount)
	driver.mu.Lock()
	cleared := driver.lastActiveInterrupt
	driver.mu.Unlock()
	require.Empty(t, cleared.threadID)
}

func TestCodexRestoreInterruptedTurnWaitFailsClosedOnSnapshotSupersede(t *testing.T) {
	var driver *CodexAppServerDriver
	driver, _ = newTestCodexDriver(t, func(req capturedRPCRequest) appServerEnvelope {
		require.Equal(t, "thread/read", req.Method)
		driver.rememberActiveInterrupt("thread-1", 2, "turn-newer")
		return appServerEnvelope{Result: rawRPCResult(t, appServerThreadResult{
			Thread: codexRecoveryThread(codexRecoveryHistory(2)),
		})}
	})
	driver.rememberActiveInterrupt("thread-1", 2, codexRecoveryInterruptTurnID)

	_, _, waited, err := driver.waitForActiveInterruptTurnForRecovery(
		context.Background(), "thread-1", 2, codexRecoveryThread(codexRecoveryHistory(2)),
	)

	require.True(t, waited)
	require.Error(t, err)
	require.Contains(t, err.Error(), "superseded")
	driver.mu.Lock()
	current := driver.lastActiveInterrupt
	driver.mu.Unlock()
	require.Equal(t, "turn-newer", current.interruptedTurnID, "old waiter must not clear the superseding interrupt")
}

func TestCodexRestoreClearsMaterializedInterruptBeforeAmbiguousRollbackRetry(t *testing.T) {
	readCount := 0
	rollbackCount := 0
	rollbackApplied := false
	driver, _ := newTestCodexDriver(t, func(req capturedRPCRequest) appServerEnvelope {
		switch req.Method {
		case "turn/interrupt":
			return appServerEnvelope{Result: rawRPCResult(t, map[string]any{})}
		case "thread/read":
			readCount++
			thread := codexRecoveryMaterializedInterruptThread()
			if rollbackApplied {
				thread = codexRecoveryThread(codexRecoveryHistory(2))
			}
			if req.Params["includeTurns"] == false {
				thread.Turns = nil
			}
			return appServerEnvelope{Result: rawRPCResult(t, appServerThreadResult{Thread: thread})}
		case "thread/rollback":
			rollbackCount++
			rollbackApplied = true
			// Model an ambiguous response: the server applied rollback, but the
			// caller did not obtain a successful result. The retry must read the
			// post-rollback count, not wait for the now-removed interrupted turn.
			return appServerEnvelope{Error: &appServerResponseError{
				Code: -32000, Message: "rollback response lost after mutation",
			}}
		default:
			t.Fatalf("unexpected method %s", req.Method)
			return appServerEnvelope{}
		}
	})
	seedCodexRecoveryActiveTurn(driver)
	require.NoError(t, driver.InterruptActiveRun(context.Background(), "thread-1"))

	err := driver.RestoreRecoveryPoint(
		context.Background(), rewindcaps.StrategyCodexAppServerTurns, codexRecoveryRefAtTwo(t),
	)
	require.Error(t, err)
	driver.mu.Lock()
	clearedBeforeRetry := driver.lastActiveInterrupt
	driver.mu.Unlock()
	require.Empty(t, clearedBeforeRetry.threadID, "materialized snapshot must clear before ambiguous rollback")

	err = driver.RestoreRecoveryPoint(
		context.Background(), rewindcaps.StrategyCodexAppServerTurns, codexRecoveryRefAtTwo(t),
	)
	require.NoError(t, err)
	require.Equal(t, 1, rollbackCount, "retry must accept the authoritative post-rollback count without a second mutation")
}