package agent

import (
	"context"
	"encoding/json"
	"testing"
	"time"

	"github.com/stretchr/testify/require"

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

func TestCodexLaunchUsesSharedProactiveCompactionThreshold(t *testing.T) {
	driver := &CodexAppServerDriver{}
	driver.resetAppServerStateForLaunchLocked(drivercontract.LaunchConfig{
		Env: map[string]string{"AGENTTEAM_PROACTIVE_COMPACTION_TOKENS": "80000"},
	}, nil)

	require.Equal(t, uint64(80_000), driver.proactiveCompactionThreshold)
}

func newCodexProactiveCompactionDriverForTest(
	t *testing.T,
	activeContextTokens uint64,
	threshold uint64,
) (*CodexAppServerDriver, *rpcCaptureWriter) {
	t.Helper()
	driver := &CodexAppServerDriver{
		launchGeneration: 7,
		codexRPCState: codexRPCState{
			pending:         make(map[int]pendingAppServerRequest),
			translatedLines: make(chan string, 8),
		},
		codexLaunchState: codexLaunchState{
			proactiveCompactionThreshold: threshold,
		},
		codexThreadState: codexThreadState{
			threadID:     "thread-1",
			threadStatus: appServerThreadStatus{Type: "idle"},
			lastTokenUsageByThread: map[string]appServerThreadTokenUsage{
				"thread-1": {
					Last: appServerTokenUsageBreakdown{
						TotalTokens: int64(activeContextTokens),
					},
				},
			},
			completedTurnIDs: make(map[string]struct{}),
		},
	}
	writer := &rpcCaptureWriter{
		t:      t,
		driver: driver,
		handle: func(capturedRPCRequest) appServerEnvelope {
			return appServerEnvelope{Result: json.RawMessage(`{}`)}
		},
	}
	driver.rawStdin = writer
	return driver, writer
}

func TestCodexProactiveIdleCompactionStartsAtSharedThreshold(t *testing.T) {
	driver, writer := newCodexProactiveCompactionDriverForTest(t, 150_000, 150_000)

	started, err := driver.StartProactiveIdleCompaction(t.Context(), "thread-1")

	require.NoError(t, err)
	require.True(t, started)
	require.Equal(t, []string{"thread/compact/start"}, capturedRPCMethods(writer.requests))
	require.Equal(t, "thread-1", writer.requests[0].Params["threadId"])
	driver.mu.Lock()
	require.True(t, driver.compactionInProgress)
	require.Equal(t, codexCompactionSourceLoopProactive, driver.compactionSource)
	require.Equal(t, appServerThreadStatus{Type: "active"}, driver.threadStatus)
	driver.mu.Unlock()
}

func TestCodexProactiveIdleCompactionSkipsWithoutDueSignal(t *testing.T) {
	tests := []struct {
		name      string
		active    uint64
		threshold uint64
		mutate    func(*CodexAppServerDriver)
	}{
		{name: "below threshold", active: 149_999, threshold: 150_000},
		{name: "disabled", active: 999_999, threshold: 0},
		{
			name:      "overflow compaction queued",
			active:    200_000,
			threshold: 150_000,
			mutate: func(driver *CodexAppServerDriver) {
				driver.compactQueued = true
			},
		},
		{
			name:      "native compaction cooldown",
			active:    200_000,
			threshold: 150_000,
			mutate: func(driver *CodexAppServerDriver) {
				driver.lastContextCompactionCompletedAt = time.Now()
			},
		},
	}
	for _, test := range tests {
		t.Run(test.name, func(t *testing.T) {
			driver, writer := newCodexProactiveCompactionDriverForTest(t, test.active, test.threshold)
			if test.mutate != nil {
				test.mutate(driver)
			}

			started, err := driver.StartProactiveIdleCompaction(t.Context(), "thread-1")

			require.NoError(t, err)
			require.False(t, started)
			require.Empty(t, writer.requests)
		})
	}
}

func TestCodexProactiveIdleCompactionUsesExactLifecycleAndSameThread(t *testing.T) {
	driver, _ := newCodexProactiveCompactionDriverForTest(t, 180_000, 150_000)
	started, err := driver.StartProactiveIdleCompaction(t.Context(), "thread-1")
	require.NoError(t, err)
	require.True(t, started)

	driver.handleNotification(appServerEnvelope{
		Method: "turn/started",
		Params: rawRPCResultForCodexCompactionTest(appServerTurnNotification{
			ThreadID: "thread-1",
			Turn: appServerTurn{
				ID:     "compact-proactive",
				Status: appServerTurnStatus{Type: "inProgress"},
			},
		}),
	})
	driver.handleNotification(codexContextCompactionItemNotificationForTest(
		"item/started", "thread-1", "compact-proactive", "compact-item",
	))
	driver.handleNotification(codexTokenUsageNotificationForTest(
		"thread-1", "compact-proactive", 0, 4_234,
	))
	driver.handleNotification(codexContextCompactionItemNotificationForTest(
		"item/completed", "thread-1", "compact-proactive", "compact-item",
	))

	driver.mu.Lock()
	require.Equal(t, codexCompactionSourceLoopProactive, driver.contextCompactionBoundary.source)
	require.False(t, driver.lastContextCompactionCompletedAt.IsZero())
	driver.mu.Unlock()

	driver.handleNotification(appServerEnvelope{
		Method: "thread/status/changed",
		Params: rawRPCResultForCodexCompactionTest(appServerThreadStatusNotification{
			ThreadID: "thread-1",
			Status:   appServerThreadStatus{Type: "idle"},
		}),
	})
	driver.handleNotification(appServerEnvelope{
		Method: "turn/completed",
		Params: rawRPCResultForCodexCompactionTest(appServerTurnNotification{
			ThreadID: "thread-1",
			Turn: appServerTurn{
				ID:     "compact-proactive",
				Status: appServerTurnStatus{Type: "completed"},
			},
		}),
	})

	compacted := driver.ParseLine(<-driver.translatedLines)
	require.Len(t, compacted, 1)
	require.Equal(t, consts.ParsedEventKindRuntimeCompacted, compacted[0].Kind)
	require.Equal(t, "thread-1", compacted[0].SessionID)
	require.Equal(t, "Context compacted (180000 → 4234 tokens).", compacted[0].Message)
	idle := driver.ParseLine(<-driver.translatedLines)
	require.Len(t, idle, 1)
	require.Equal(t, consts.ParsedEventKindRuntimeIdle, idle[0].Kind)
	require.Equal(t, "thread-1", idle[0].SessionID)
	driver.mu.Lock()
	require.False(t, driver.compactionInProgress)
	require.Equal(t, appServerThreadStatus{Type: "idle"}, driver.threadStatus)
	driver.mu.Unlock()
}

func TestCodexProactiveIdleCompactionTerminalFailureIsMaintenanceOnly(t *testing.T) {
	tests := []struct {
		name string
		fail func(*CodexAppServerDriver)
	}{
		{
			name: "error notification",
			fail: func(driver *CodexAppServerDriver) {
				driver.handleNotification(appServerEnvelope{
					Method: consts.ParsedEventKindError,
					Params: rawRPCResultForCodexCompactionTest(appServerErrorNotification{
						ThreadID:  "thread-1",
						TurnID:    "compact-proactive",
						WillRetry: false,
						Error:     json.RawMessage(`{"message":"summary failed","codexErrorInfo":"contextWindowExceeded"}`),
					}),
				})
			},
		},
		{
			name: "failed completion without error notification",
			fail: func(driver *CodexAppServerDriver) {
				driver.handleNotification(appServerEnvelope{
					Method: "turn/completed",
					Params: rawRPCResultForCodexCompactionTest(appServerTurnNotification{
						ThreadID: "thread-1",
						Turn: appServerTurn{
							ID:     "compact-proactive",
							Status: appServerTurnStatus{Type: "failed"},
						},
					}),
				})
			},
		},
	}
	for _, test := range tests {
		t.Run(test.name, func(t *testing.T) {
			driver, _ := newCodexProactiveCompactionDriverForTest(t, 180_000, 150_000)
			started, err := driver.StartProactiveIdleCompaction(t.Context(), "thread-1")
			require.NoError(t, err)
			require.True(t, started)
			driver.handleNotification(appServerEnvelope{
				Method: "turn/started",
				Params: rawRPCResultForCodexCompactionTest(appServerTurnNotification{
					ThreadID: "thread-1",
					Turn: appServerTurn{
						ID:     "compact-proactive",
						Status: appServerTurnStatus{Type: "inProgress"},
					},
				}),
			})
			driver.handleNotification(appServerEnvelope{
				Method: "thread/status/changed",
				Params: rawRPCResultForCodexCompactionTest(appServerThreadStatusNotification{
					ThreadID: "thread-1",
					Status:   appServerThreadStatus{Type: "idle"},
				}),
			})

			test.fail(driver)

			errorEvents := driver.ParseLine(<-driver.translatedLines)
			require.Len(t, errorEvents, 1)
			require.Equal(t, consts.ParsedEventKindError, errorEvents[0].Kind)
			require.True(t, parsedBoolMetadata(errorEvents[0].Metadata, contextCompactionFailedMetadataKey))
			require.True(t, parsedBoolMetadata(errorEvents[0].Metadata, contextCompactionMaintenanceMetadataKey))
			idleEvents := driver.ParseLine(<-driver.translatedLines)
			require.Len(t, idleEvents, 1)
			require.Equal(t, consts.ParsedEventKindRuntimeIdle, idleEvents[0].Kind)
			require.Equal(t, "thread-1", idleEvents[0].SessionID)
			driver.mu.Lock()
			require.False(t, driver.compactQueued, "maintenance failure must not queue overflow recovery")
			driver.mu.Unlock()
		})
	}
}

func TestCodexProactiveIdleCompactionDefiniteRejectionRestoresIdle(t *testing.T) {
	writer := &codexTransportSignalWriter{wrote: make(chan struct{})}
	driver, _ := newCodexProactiveCompactionDriverForTest(t, 180_000, 150_000)
	driver.rawStdin = writer
	ctx, cancel := context.WithCancel(t.Context())
	cancel()

	started, err := driver.StartProactiveIdleCompaction(ctx, "thread-1")

	require.False(t, started)
	require.ErrorIs(t, err, context.Canceled)
	require.NotErrorIs(t, err, errCodexAppServerTransportAmbiguous)
	driver.mu.Lock()
	require.False(t, driver.compactionInProgress)
	require.Empty(t, driver.compactionSource)
	require.Equal(t, appServerThreadStatus{Type: "idle"}, driver.threadStatus)
	driver.mu.Unlock()
}

func TestCodexProactiveIdleCompactionAmbiguousStartStaysActive(t *testing.T) {
	writer := &codexTransportSignalWriter{wrote: make(chan struct{})}
	driver, _ := newCodexProactiveCompactionDriverForTest(t, 180_000, 150_000)
	driver.rawStdin = writer
	ctx, cancel := context.WithCancel(t.Context())
	type startResult struct {
		started bool
		err     error
	}
	result := make(chan startResult, 1)
	go func() {
		started, err := driver.StartProactiveIdleCompaction(ctx, "thread-1")
		result <- startResult{started: started, err: err}
	}()
	<-writer.wrote
	cancel()

	got := <-result
	require.True(t, got.started)
	require.ErrorIs(t, got.err, errCodexAppServerTransportAmbiguous)
	driver.mu.Lock()
	require.True(t, driver.compactionInProgress)
	require.Equal(t, codexCompactionSourceLoopProactive, driver.compactionSource)
	require.Equal(t, appServerThreadStatus{Type: "active"}, driver.threadStatus)
	driver.mu.Unlock()
}

func TestCodexProactiveIdleCompactionRPCRejectionAfterTurnStartIsAmbiguous(t *testing.T) {
	driver, writer := newCodexProactiveCompactionDriverForTest(t, 180_000, 150_000)
	writer.handle = func(capturedRPCRequest) appServerEnvelope {
		driver.handleNotification(appServerEnvelope{
			Method: "turn/started",
			Params: rawRPCResultForCodexCompactionTest(appServerTurnNotification{
				ThreadID: "thread-1",
				Turn: appServerTurn{
					ID:     "compact-proactive",
					Status: appServerTurnStatus{Type: "inProgress"},
				},
			}),
		})
		return appServerEnvelope{
			Error: &appServerResponseError{Code: -32000, Message: "request rejected"},
		}
	}

	started, err := driver.StartProactiveIdleCompaction(t.Context(), "thread-1")

	require.True(t, started)
	require.Error(t, err)
	driver.mu.Lock()
	require.True(t, driver.compactionInProgress)
	require.Equal(t, "compact-proactive", driver.compactionTurnID)
	require.Equal(t, codexCompactionSourceLoopProactive, driver.compactionSource)
	driver.mu.Unlock()
}