package agent

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

	"github.com/stretchr/testify/require"

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

func newActiveCodexCompactionDriverForTest() *CodexAppServerDriver {
	driver := &CodexAppServerDriver{
		launchGeneration: 2,
		codexRPCState: codexRPCState{
			pending:         make(map[int]pendingAppServerRequest),
			translatedLines: make(chan string, 8),
		},
		codexThreadState: codexThreadState{
			threadID:         "thread-1",
			threadStatus:     appServerThreadStatus{Type: "idle"},
			completedTurnIDs: make(map[string]struct{}),
		},
	}
	driver.beginCompaction("thread-1")
	driver.handleNotification(appServerEnvelope{
		Method: "turn/started",
		Params: rawRPCResultForCodexCompactionTest(appServerTurnNotification{
			ThreadID: "thread-1",
			Turn: appServerTurn{
				ID:     "compact-1",
				Status: appServerTurnStatus{Type: "inProgress"},
			},
		}),
	})
	return driver
}

func rawRPCResultForCodexCompactionTest(value any) json.RawMessage {
	data, err := json.Marshal(value)
	if err != nil {
		panic(err)
	}
	return data
}

func codexCompactionStatusNotificationForTest() appServerEnvelope {
	return appServerEnvelope{
		Method: "thread/status/changed",
		Params: rawRPCResultForCodexCompactionTest(appServerThreadStatusNotification{
			ThreadID: "thread-1",
			Status:   appServerThreadStatus{Type: "idle"},
		}),
	}
}

func codexCompactionCompletedNotificationForTest(turnID string) appServerEnvelope {
	return appServerEnvelope{
		Method: "turn/completed",
		Params: rawRPCResultForCodexCompactionTest(appServerTurnNotification{
			ThreadID: "thread-1",
			Turn: appServerTurn{
				ID:     turnID,
				Status: appServerTurnStatus{Type: "completed"},
			},
		}),
	}
}

func codexContextCompactionItemNotificationForTest(method, threadID, turnID, itemID string) appServerEnvelope {
	return appServerEnvelope{
		Method: method,
		Params: rawRPCResultForCodexCompactionTest(appServerItemNotification{
			ThreadID: threadID,
			TurnID:   turnID,
			Item: appServerItem{
				Type: "contextCompaction",
				ID:   itemID,
			},
		}),
	}
}

func codexTokenUsageNotificationForTest(
	threadID string,
	turnID string,
	inputTokens int64,
	totalTokens int64,
) appServerEnvelope {
	contextWindow := int64(353_400)
	return appServerEnvelope{
		Method: "thread/tokenUsage/updated",
		Params: rawRPCResultForCodexCompactionTest(appServerThreadTokenUsageNotification{
			ThreadID: threadID,
			TurnID:   turnID,
			TokenUsage: appServerThreadTokenUsage{
				ModelContextWindow: &contextWindow,
				Last: appServerTokenUsageBreakdown{
					InputTokens: inputTokens,
					TotalTokens: totalTokens,
				},
				Total: appServerTokenUsageBreakdown{
					InputTokens: inputTokens,
					TotalTokens: totalTokens,
				},
			},
		}),
	}
}

func parseNextCodexEventForCompactionTest(t *testing.T, driver *CodexAppServerDriver) drivercontract.ParsedEvent {
	t.Helper()
	events := driver.ParseLine(<-driver.translatedLines)
	require.Len(t, events, 1)
	return events[0]
}

func newCodexCompactionObservationDriver(threadID, turnID, model string) *CodexAppServerDriver {
	status := appServerThreadStatus{Type: "idle"}
	if turnID != "" {
		status = appServerThreadStatus{Type: "active"}
	}
	modelByThread := make(map[string]string)
	if model != "" {
		modelByThread[threadID] = model
	}
	return &CodexAppServerDriver{
		codexRPCState: codexRPCState{
			translatedLines: make(chan string, 8),
		},
		codexLaunchState: codexLaunchState{effectiveModel: model},
		codexThreadState: codexThreadState{
			threadID:                  threadID,
			threadStatus:              status,
			activeTurnID:              turnID,
			completedTurnIDs:          make(map[string]struct{}),
			streamedItemIDs:           make(map[string]struct{}),
			lastTokenUsageByThread:    make(map[string]appServerThreadTokenUsage),
			pendingTokenUsageByThread: make(map[string][]pendingAppServerThreadTokenUsage),
			modelByThread:             modelByThread,
			modelByThreadTurn:         make(map[string]string),
			modelHydrationInFlight:    make(map[string]uint64),
			modelHydrationBlocked:     make(map[string]struct{}),
		},
	}
}

func TestCodexAppServerNativeContextCompactionMarksOnlyFirstPostBoundaryRootUsage(t *testing.T) {
	driver := newCodexCompactionObservationDriver("thread-root", "turn-root", "gpt-5.4")
	driver.launchGeneration = 3
	driver.modelByThread["thread-child"] = "gpt-5-mini"

	driver.handleNotification(codexTokenUsageNotificationForTest("thread-root", "turn-root", 330_000, 338_000))
	require.False(t, parseNextCodexEventForCompactionTest(t, driver).ContextCompacted)

	driver.handleNotification(codexContextCompactionItemNotificationForTest(
		"item/started", "thread-root", "turn-root", "compact-item",
	))
	// Codex recomputes active context after installing compacted history with
	// last.totalTokens populated and every billable-token component zero.
	driver.handleNotification(codexTokenUsageNotificationForTest("thread-root", "turn-root", 0, 10_500))
	require.Empty(t, driver.translatedLines, "recomputed active context is not a model invocation")
	driver.handleNotification(codexContextCompactionItemNotificationForTest(
		"item/completed", "thread-root", "turn-root", "compact-item",
	))
	driver.handleNotification(codexContextCompactionItemNotificationForTest(
		"item/started", "thread-root", "turn-root", "compact-item",
	))
	driver.handleNotification(codexContextCompactionItemNotificationForTest(
		"item/completed", "thread-root", "turn-root", "compact-item",
	))
	require.Empty(t, driver.translatedLines, "compaction diagnostics must not create trajectory text")

	driver.mu.Lock()
	boundary := driver.contextCompactionBoundary
	driver.mu.Unlock()
	require.Equal(t, codexCompactionSourceRuntimeNative, boundary.source)
	require.True(t, appServerThreadTokenUsageKnown(boundary.before))
	require.Equal(t, uint64(338_000), nonNegativeInt64TokenCount(boundary.before.Last.TotalTokens))
	require.True(t, appServerThreadTokenUsageKnown(boundary.after))
	require.Equal(t, uint64(10_500), nonNegativeInt64TokenCount(boundary.after.Last.TotalTokens))

	// Foreign subagent usage remains ordinary and cannot consume the root latch.
	driver.handleNotification(codexTokenUsageNotificationForTest("thread-child", "turn-child", 2_000, 2_500))
	childUsage := parseNextCodexEventForCompactionTest(t, driver)
	require.Equal(t, "thread-child", childUsage.SessionID)
	require.False(t, childUsage.ContextCompacted)

	driver.handleNotification(codexTokenUsageNotificationForTest("thread-root", "turn-root", 9_000, 9_500))
	firstRootUsage := parseNextCodexEventForCompactionTest(t, driver)
	require.True(t, firstRootUsage.ContextCompacted)
	driver.handleNotification(codexTokenUsageNotificationForTest("thread-root", "turn-root", 8_000, 8_500))
	require.False(t, parseNextCodexEventForCompactionTest(t, driver).ContextCompacted)
}

func TestCodexAppServerEmergencyContextCompactionCarriesExactOverflowCause(t *testing.T) {
	driver := newCodexCompactionObservationDriver("thread-root", "", "")
	driver.launchGeneration = 4
	driver.queueCompaction("thread-root", "overflow-turn")
	threadID, ok := driver.prepareIdleCompaction("thread-root")
	require.True(t, ok)
	driver.beginCompaction(threadID)
	driver.handleNotification(appServerEnvelope{
		Method: "turn/started",
		Params: rawRPCResultForCodexCompactionTest(appServerTurnNotification{
			ThreadID: "thread-root",
			Turn: appServerTurn{
				ID:     "emergency-compact-turn",
				Status: appServerTurnStatus{Type: "inProgress"},
			},
		}),
	})
	driver.handleNotification(codexContextCompactionItemNotificationForTest(
		"item/started", "thread-root", "emergency-compact-turn", "compact-item",
	))
	driver.handleNotification(codexContextCompactionItemNotificationForTest(
		"item/completed", "thread-root", "emergency-compact-turn", "compact-item",
	))

	driver.mu.Lock()
	boundary := driver.contextCompactionBoundary
	driver.mu.Unlock()
	require.Equal(t, codexCompactionSourceAgentTeamEmergency, boundary.source)
	require.Equal(t, "overflow-turn", boundary.causeTurnID)
	require.Equal(t, "emergency-compact-turn", boundary.turnID)
	require.Equal(t, "compact-item", boundary.itemID)
	require.Empty(t, driver.translatedLines)
}

func TestCodexAppServerConflictingContextCompactionItemPairIsUnknown(t *testing.T) {
	driver := newCodexCompactionObservationDriver("thread-root", "turn-root", "")
	driver.handleNotification(codexContextCompactionItemNotificationForTest(
		"item/started", "thread-root", "turn-root", "compact-a",
	))
	driver.handleNotification(codexContextCompactionItemNotificationForTest(
		"item/started", "thread-root", "turn-root", "compact-b",
	))
	driver.handleNotification(codexContextCompactionItemNotificationForTest(
		"item/completed", "thread-root", "turn-root", "compact-b",
	))

	driver.mu.Lock()
	boundary := driver.contextCompactionBoundary
	driver.mu.Unlock()
	require.Equal(t, codexCompactionSourceUnknown, boundary.source)
	require.Equal(t, "compact-b", boundary.itemID)
	require.Empty(t, driver.translatedLines)
}

func TestCodexAppServerPostCompactionUsageRetainsBoundaryWhileModelHydrates(t *testing.T) {
	driver := newCodexCompactionObservationDriver("thread-root", "turn-root", "")
	driver.handleNotification(codexContextCompactionItemNotificationForTest(
		"item/started", "thread-root", "turn-root", "compact-item",
	))
	driver.handleNotification(codexContextCompactionItemNotificationForTest(
		"item/completed", "thread-root", "turn-root", "compact-item",
	))
	driver.handleNotification(codexTokenUsageNotificationForTest("thread-root", "turn-root", 9_000, 9_500))
	require.Empty(t, driver.translatedLines)

	driver.mu.Lock()
	require.Len(t, driver.pendingTokenUsageByThread["thread-root"], 1)
	require.True(t, driver.pendingTokenUsageByThread["thread-root"][0].contextCompacted)
	driver.mu.Unlock()

	driver.updateEffectiveModel("gpt-5.4")
	usage := parseNextCodexEventForCompactionTest(t, driver)
	require.Equal(t, "thread-root", usage.SessionID)
	require.Equal(t, "gpt-5.4", usage.Model)
	require.True(t, usage.ContextCompacted)
}

func TestCodexAppServerContextCompactionRequiresExactRootUsageIdentity(t *testing.T) {
	driver := newCodexCompactionObservationDriver("thread-root", "turn-root", "gpt-5.4")
	driver.handleNotification(codexContextCompactionItemNotificationForTest(
		"item/started", "thread-root", "turn-root", "compact-item",
	))
	driver.handleNotification(codexContextCompactionItemNotificationForTest(
		"item/completed", "thread-root", "turn-root", "compact-item",
	))

	driver.handleNotification(codexTokenUsageNotificationForTest("", "turn-root", 9_000, 9_500))
	require.False(t, parseNextCodexEventForCompactionTest(t, driver).ContextCompacted)
	driver.handleNotification(codexTokenUsageNotificationForTest("thread-root", "turn-root", 8_000, 8_500))
	require.True(t, parseNextCodexEventForCompactionTest(t, driver).ContextCompacted)
}

func TestCodexAppServerTerminalContextOverflowCarriesCorrelation(t *testing.T) {
	driver := newCodexCompactionObservationDriver("thread-root", "turn-overflow", "gpt-5.4")
	driver.handleNotification(appServerEnvelope{
		Method: consts.ParsedEventKindError,
		Params: rawRPCResultForCodexCompactionTest(map[string]any{
			"threadId":  "thread-root",
			"turnId":    "turn-overflow",
			"willRetry": false,
			"error": map[string]any{
				"message":        "request exceeded the context window",
				"codexErrorInfo": "contextWindowExceeded",
			},
		}),
	})

	errorEvent := parseNextCodexEventForCompactionTest(t, driver)
	require.True(t, parsedBoolMetadata(errorEvent.Metadata, "context_window_exceeded"))
	require.Equal(t, "thread-root", parsedStringMetadata(errorEvent.Metadata, codexThreadIDMetadataKey))
	require.Equal(t, "turn-overflow", parsedStringMetadata(errorEvent.Metadata, codexTurnIDMetadataKey))
	require.Equal(t, consts.ParsedEventKindTurnEnd, parseNextCodexEventForCompactionTest(t, driver).Kind)
	driver.mu.Lock()
	require.True(t, driver.compactQueued)
	require.Equal(t, "turn-overflow", driver.queuedCompactionCauseTurnID)
	driver.mu.Unlock()
}

func TestCodexAppServerCompactionConsumesUpstreamIdleBeforeMatchingCompletion(t *testing.T) {
	driver := newActiveCodexCompactionDriverForTest()

	// Codex app-server publishes thread/status/changed idle before it publishes
	// turn/completed for the same turn. The early status is not a boundary until
	// the concrete compaction id completes.
	driver.handleNotification(codexCompactionStatusNotificationForTest())
	driver.handleNotification(codexCompactionStatusNotificationForTest())
	require.Empty(t, driver.translatedLines)
	driver.mu.Lock()
	require.True(t, driver.compactionInProgress)
	require.Equal(t, "compact-1", driver.compactionTurnID)
	require.Equal(t, "idle", driver.compactionTerminalStatus.Type)
	driver.mu.Unlock()

	// A delayed completion from another turn must not consume the latched status.
	driver.handleNotification(codexCompactionCompletedNotificationForTest("turn-stale"))
	require.Empty(t, driver.translatedLines)
	driver.mu.Lock()
	require.True(t, driver.compactionInProgress)
	require.Equal(t, "compact-1", driver.compactionTurnID)
	require.Equal(t, "idle", driver.compactionTerminalStatus.Type)
	driver.mu.Unlock()

	driver.handleNotification(codexCompactionCompletedNotificationForTest("compact-1"))
	compacted := parseNextCodexEventForCompactionTest(t, driver)
	require.Equal(t, consts.ParsedEventKindRuntimeCompacted, compacted.Kind)
	require.Equal(t, "thread-1", compacted.SessionID)
	idle := parseNextCodexEventForCompactionTest(t, driver)
	require.Equal(t, consts.ParsedEventKindRuntimeIdle, idle.Kind)
	require.Equal(t, "thread-1", idle.SessionID)
	driver.mu.Lock()
	require.False(t, driver.compactionInProgress)
	require.Empty(t, driver.compactionTurnID)
	require.Empty(t, driver.compactionTerminalStatus.Type)
	require.Equal(t, "idle", driver.threadStatus.Type)
	driver.mu.Unlock()

	// Both duplicate terminal signals are idempotent.
	driver.handleNotification(codexCompactionCompletedNotificationForTest("compact-1"))
	driver.handleNotification(codexCompactionStatusNotificationForTest())
	require.Empty(t, driver.translatedLines)
}

func TestCodexAppServerOldGenerationCannotConsumeCompactionTerminalLatch(t *testing.T) {
	driver := newActiveCodexCompactionDriverForTest()

	for _, notification := range []appServerEnvelope{
		codexCompactionStatusNotificationForTest(),
		codexCompactionCompletedNotificationForTest("compact-1"),
	} {
		frame, err := json.Marshal(notification)
		require.NoError(t, err)
		require.ErrorIs(t, driver.handleAppServerLineForGeneration(string(frame), 1), errCodexAppServerTransportSuperseded)
	}
	require.Empty(t, driver.translatedLines)
	driver.mu.Lock()
	require.True(t, driver.compactionInProgress)
	require.Equal(t, "compact-1", driver.compactionTurnID)
	require.Empty(t, driver.compactionTerminalStatus.Type)
	_, staleCompletionRecorded := driver.completedTurnIDs["compact-1"]
	driver.mu.Unlock()
	require.False(t, staleCompletionRecorded)

	for _, notification := range []appServerEnvelope{
		codexCompactionStatusNotificationForTest(),
		codexCompactionCompletedNotificationForTest("compact-1"),
	} {
		frame, err := json.Marshal(notification)
		require.NoError(t, err)
		require.NoError(t, driver.handleAppServerLineForGeneration(string(frame), 2))
	}
	require.Equal(t,
		consts.ParsedEventKindRuntimeCompacted,
		parseNextCodexEventForCompactionTest(t, driver).Kind)
	require.Equal(t,
		consts.ParsedEventKindRuntimeIdle,
		parseNextCodexEventForCompactionTest(t, driver).Kind)
}

func TestCodexAppServerOrdinaryTurnDoesNotUseCompactionTerminalLatch(t *testing.T) {
	driver := &CodexAppServerDriver{
		codexRPCState: codexRPCState{translatedLines: make(chan string, 4)},
		codexThreadState: codexThreadState{
			threadID:         "thread-1",
			threadStatus:     appServerThreadStatus{Type: "active"},
			activeTurnID:     "turn-1",
			completedTurnIDs: make(map[string]struct{}),
		},
	}

	// Codex can publish idle immediately before the ordinary turn/completed too.
	// It must not clear the concrete active turn or create the compaction-only
	// latch; the authoritative ordinary boundary remains TurnEnd.
	driver.handleNotification(codexCompactionStatusNotificationForTest())
	require.Empty(t, driver.translatedLines)
	driver.mu.Lock()
	require.Equal(t, "active", driver.threadStatus.Type)
	require.Equal(t, "turn-1", driver.activeTurnID)
	require.Empty(t, driver.compactionTerminalStatus.Type)
	driver.mu.Unlock()

	driver.handleNotification(codexCompactionCompletedNotificationForTest("turn-1"))
	line := <-driver.translatedLines
	events := driver.ParseLine(line)
	require.Len(t, events, 1)
	require.Equal(t, consts.ParsedEventKindTurnEnd, events[0].Kind)
	driver.mu.Lock()
	require.Equal(t, "idle", driver.threadStatus.Type)
	require.Empty(t, driver.activeTurnID)
	require.False(t, driver.compactionInProgress)
	require.Empty(t, driver.compactionTerminalStatus.Type)
	driver.mu.Unlock()

	driver.handleNotification(codexCompactionCompletedNotificationForTest("turn-1"))
	driver.handleNotification(codexCompactionStatusNotificationForTest())
	require.Empty(t, driver.translatedLines)
}

func TestCodexAppServerAuthoritativeReadPreservesExactActiveCompactionFence(t *testing.T) {
	readResult := appServerThreadResult{Thread: appServerThread{
		ID:     "thread-1",
		Status: appServerThreadStatus{Type: "active"},
		Turns: []appServerTurn{{
			ID:     "compact-1",
			Status: appServerTurnStatus{Type: "inProgress"},
		}},
	}}
	driver, writer := newTestCodexDriver(t, func(req capturedRPCRequest) appServerEnvelope {
		require.Equal(t, "thread/read", req.Method, "delivery must not steer into the compaction turn")
		require.Equal(t, "thread-1", req.Params["threadId"])
		require.Equal(t, true, req.Params["includeTurns"])
		return appServerEnvelope{Result: rawRPCResultForCodexCompactionTest(readResult)}
	})
	driver.translatedLines = make(chan string, 4)
	driver.threadID = "thread-1"
	driver.threadStatus = appServerThreadStatus{Type: "idle"}
	driver.completedTurnIDs = make(map[string]struct{})
	driver.beginCompaction("thread-1")
	driver.handleNotification(appServerEnvelope{
		Method: "turn/started",
		Params: rawRPCResultForCodexCompactionTest(appServerTurnNotification{
			ThreadID: "thread-1",
			Turn: appServerTurn{
				ID:     "compact-1",
				Status: appServerTurnStatus{Type: "inProgress"},
			},
		}),
	})

	// An authoritative read can occur during reconciliation/recovery while the
	// compaction is active. The same exact in-progress turn confirms the cached
	// classification; it must not turn into a regular steerable active turn.
	require.NoError(t, driver.refreshLiveDeliverySnapshot(context.Background(), "thread-1"))
	driver.mu.Lock()
	require.True(t, driver.compactionInProgress)
	require.Equal(t, "compact-1", driver.compactionTurnID)
	require.Equal(t, "compact-1", driver.activeTurnID)
	driver.mu.Unlock()

	err := driver.DeliverAckedBusyRuntimeInput(context.Background(), testRuntimeInputWithConfig(drivercontract.RuntimeInput{ID: "msg-during-compaction"}, "must not steer"), "thread-1")
	require.Error(t, err)
	require.Contains(t, err.Error(), "compaction")
	require.Equal(t, []string{"thread/read"}, capturedRPCMethods(writer.requests))

	// Codex publishes idle before turn/completed. Even a second active snapshot
	// of the same turn must retain that latch until the matching completion.
	driver.handleNotification(codexCompactionStatusNotificationForTest())
	require.NoError(t, driver.refreshLiveDeliverySnapshot(context.Background(), "thread-1"))
	driver.mu.Lock()
	require.True(t, driver.compactionInProgress)
	require.Equal(t, "compact-1", driver.compactionTurnID)
	require.Equal(t, "idle", driver.compactionTerminalStatus.Type)
	driver.mu.Unlock()

	driver.handleNotification(codexCompactionCompletedNotificationForTest("compact-1"))
	require.Equal(t,
		consts.ParsedEventKindRuntimeCompacted,
		parseNextCodexEventForCompactionTest(t, driver).Kind)
	require.Equal(t,
		consts.ParsedEventKindRuntimeIdle,
		parseNextCodexEventForCompactionTest(t, driver).Kind)
	require.Equal(t, []string{"thread/read", "thread/read"}, capturedRPCMethods(writer.requests))
}

func TestCodexAppServerAuthoritativeSnapshotClearsReplacedOrTerminalCompaction(t *testing.T) {
	for _, test := range []struct {
		name     string
		snapshot appServerThread
	}{
		{
			name: "different active turn",
			snapshot: appServerThread{
				ID:     "thread-1",
				Status: appServerThreadStatus{Type: "active"},
				Turns: []appServerTurn{{
					ID:     "replacement-turn",
					Status: appServerTurnStatus{Type: "inProgress"},
				}},
			},
		},
		{
			name: "terminal snapshot",
			snapshot: appServerThread{
				ID:     "thread-1",
				Status: appServerThreadStatus{Type: "idle"},
				Turns: []appServerTurn{{
					ID:     "compact-1",
					Status: appServerTurnStatus{Type: "completed"},
				}},
			},
		},
	} {
		t.Run(test.name, func(t *testing.T) {
			driver := newActiveCodexCompactionDriverForTest()
			driver.mu.Lock()
			driver.compactionTerminalStatus = appServerThreadStatus{Type: "idle"}
			driver.mu.Unlock()

			driver.updateThreadSnapshot(test.snapshot)

			driver.mu.Lock()
			require.False(t, driver.compactionInProgress)
			require.Empty(t, driver.compactionTurnID)
			require.Empty(t, driver.compactionTerminalStatus.Type)
			driver.mu.Unlock()
		})
	}
}