package agent

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

	"github.com/stretchr/testify/require"

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

type promptAttributionRecorderStub struct {
	compile  []metrics.PromptCompileRecord
	usage    []metrics.UsageObservationRecord
	exposure []metrics.InvocationExposureRecord
}

func (r *promptAttributionRecorderStub) RecordPromptCompileMetrics(items []metrics.PromptCompileRecord) {
	r.compile = append(r.compile, items...)
}

func (r *promptAttributionRecorderStub) RecordUsageObservation(
	item metrics.UsageObservationRecord,
	exposure *metrics.InvocationExposureRecord,
) {
	r.usage = append(r.usage, item)
	if exposure != nil {
		r.exposure = append(r.exposure, *exposure)
	}
}

func TestRecordUsageKeepsInvocationIdentityRunDimensionsAndExposure(t *testing.T) {
	recorder := &promptAttributionRecorderStub{}
	previous := tokenAnalysisRecorder
	tokenAnalysisRecorder = func() metrics.TokenAnalysisRecorder { return recorder }
	t.Cleanup(func() { tokenAnalysisRecorder = previous })
	process := &Process{
		AgentID:   "agent-1",
		SessionID: "thread-primary",
		processRuntimeBinding: processRuntimeBinding{
			driver: &CodexAppServerDriver{},
		},
		processUsageState: processUsageState{
			model:           "configured-model",
			launchID:        "launch-1",
			launchStartedAt: time.UnixMilli(1_700_000_000_000),
			runtime:         "codex",
			sessionKind:     "interactive",
			providerAPIKey:  "",
			activePromptMetrics: map[string]metrics.InvocationExposureBlock{
				"platform.agentteam.core\x00persistent": {
					BlockID: "platform.agentteam.core", Durability: "persistent", Tokens: 42,
				},
			},
			contextExposureComplete: true,
		},
	}

	process.recordUsage(drivercontract.ParsedEvent{
		SessionID:            "thread-primary",
		Model:                "effective-model",
		ModelInvocationID:    "turn-1",
		ContextExposureKnown: true,
		InputTokens:          100,
		CachedInputTokens:    20,
		OutputTokens:         10,
	})

	require.Equal(t, []metrics.UsageObservationRecord{{
		LaunchStartedAt: 1_700_000_000_000,
		AgentID:         "agent-1",
		Runtime:         "codex",
		LaunchID:        "launch-1",
		SessionKind:     "interactive",
		SessionID:       "thread-primary",
		Model:           "effective-model",
		InvocationID:    "turn-1",
		InvocationCause: "unknown",
		InputTokens:     100,
		CachedTokens:    20,
		OutputTokens:    10,
	}}, recorder.usage)
	require.Equal(t, []metrics.InvocationExposureRecord{{
		Complete: true,
		Blocks: []metrics.InvocationExposureBlock{{
			BlockID: "platform.agentteam.core", Durability: "persistent", Tokens: 42,
		}},
	}}, recorder.exposure)
}

func TestRecordUsageKeepsLowCardinalityTotalsWhenModelIsAmbiguous(t *testing.T) {
	require.Equal(t, "unknown/claude", unknownUsageModel(" claude "))
	require.Equal(t, "unknown", unknownUsageModel(""))
}

func TestOpenCodeInvocationExposureWaitsForCompactionBoundary(t *testing.T) {
	recorder := &promptAttributionRecorderStub{}
	previous := tokenAnalysisRecorder
	tokenAnalysisRecorder = func() metrics.TokenAnalysisRecorder { return recorder }
	t.Cleanup(func() { tokenAnalysisRecorder = previous })
	ap := &Process{
		AgentID: "agent-1",
		processUsageState: processUsageState{
			launchID:                "launch-1",
			launchStartedAt:         time.UnixMilli(1_700_000_000_000),
			runtime:                 "opencode",
			sessionKind:             "interactive",
			contextExposureComplete: true,
			activePromptMetrics: map[string]metrics.InvocationExposureBlock{
				"core\x00persistent":       {BlockID: "core", Durability: "persistent", Tokens: 40},
				"reference\x00compactable": {BlockID: "reference", Durability: "compactable", Tokens: 20},
			},
		},
	}
	event := drivercontract.ParsedEvent{
		SessionID:            "session-1",
		ModelInvocationID:    "step-1",
		ContextExposureKnown: true,
		ContextCompacted:     true,
		InputTokens:          100,
	}
	ap.recordUsage(event)
	require.Len(t, recorder.exposure, 1)
	require.Equal(t, []metrics.InvocationExposureBlock{{
		BlockID: "core", Durability: "persistent", Tokens: 40,
	}}, recorder.exposure[0].Blocks)
}

func TestCompactionPrunesActiveExposureWhenCollectionIsDisabled(t *testing.T) {
	previous := tokenAnalysisRecorder
	tokenAnalysisRecorder = func() metrics.TokenAnalysisRecorder { return nil }
	t.Cleanup(func() { tokenAnalysisRecorder = previous })
	ap := &Process{
		processUsageState: processUsageState{
			activePromptMetrics: map[string]metrics.InvocationExposureBlock{
				"core\x00persistent":       {BlockID: "core", Durability: "persistent", Tokens: 40},
				"reference\x00compactable": {BlockID: "reference", Durability: "compactable", Tokens: 20},
			},
		},
	}

	ap.recordUsage(drivercontract.ParsedEvent{
		ModelInvocationID:    "step-1",
		ContextExposureKnown: true,
		ContextCompacted:     true,
		InputTokens:          100,
	})

	require.Equal(t, map[string]metrics.InvocationExposureBlock{
		"core\x00persistent": {BlockID: "core", Durability: "persistent", Tokens: 40},
	}, ap.activePromptMetrics)
}

func TestRecordUsageExcludesRuntimeNativeInvocationFromRootExposure(t *testing.T) {
	recorder := &promptAttributionRecorderStub{}
	previous := tokenAnalysisRecorder
	tokenAnalysisRecorder = func() metrics.TokenAnalysisRecorder { return recorder }
	t.Cleanup(func() { tokenAnalysisRecorder = previous })
	ap := &Process{
		AgentID: "agent-1",
		processUsageState: processUsageState{
			launchID:                "launch-1",
			runtime:                 "kimi",
			sessionKind:             "interactive",
			contextExposureComplete: true,
			activePromptMetrics: map[string]metrics.InvocationExposureBlock{
				"core\x00persistent": {BlockID: "core", Durability: "persistent", Tokens: 40},
			},
		},
	}

	ap.recordUsage(drivercontract.ParsedEvent{
		SessionID:               "session-1",
		ModelInvocationID:       "subagent-step-1",
		ContextExposureExcluded: true,
		InputTokens:             100,
	})

	require.Len(t, recorder.usage, 1)
	require.True(t, recorder.usage[0].ExposureExcluded)
	require.Empty(t, recorder.exposure)
}

func TestActivatePromptBoundaryReplacesBoundaryScopedBlocks(t *testing.T) {
	ap := &Process{processUsageState: processUsageState{
		contextExposureComplete: true,
		promptMetricsBaselineOK: true,
		activePromptMetrics: map[string]metrics.InvocationExposureBlock{
			"core\x00persistent":               {BlockID: "core", Durability: "persistent", Tokens: 40},
			"old-turn\x00turn_only":            {BlockID: "old-turn", Durability: "turn_only", Tokens: 10},
			"old-notice\x00until_acknowledged": {BlockID: "old-notice", Durability: "until_acknowledged", Tokens: 5},
		},
	}}

	ap.activatePromptBoundary(promptMetricsBoundary{
		rendered: promptcompiler.RenderedPromptPayload{
			AttributionBlocks: []promptcompiler.PromptBlock{
				{ID: "new-turn", Durability: promptcompiler.DurabilityTurnOnly, TokenEstimate: 11},
				{ID: "new-notice", Durability: promptcompiler.DurabilityUntilAcknowledged, TokenEstimate: 6},
			},
		},
	})

	require.Equal(t, map[string]metrics.InvocationExposureBlock{
		"core\x00persistent":               {BlockID: "core", Durability: "persistent", Tokens: 40},
		"new-turn\x00turn_only":            {BlockID: "new-turn", Durability: "turn_only", Tokens: 11},
		"old-notice\x00until_acknowledged": {BlockID: "old-notice", Durability: "until_acknowledged", Tokens: 5},
		"new-notice\x00until_acknowledged": {BlockID: "new-notice", Durability: "until_acknowledged", Tokens: 6},
	}, ap.activePromptMetrics)
}

func TestRetainedLaunchPromptMetricsFlushOnceWhenCollectionTurnsOn(t *testing.T) {
	recorder := &promptAttributionRecorderStub{}
	previous := promptAttributionRecorder
	promptAttributionRecorder = func() metrics.PromptAttributionRecorder {
		if !metrics.DetailedTokenCollectionEnabled() {
			return nil
		}
		return recorder
	}
	t.Cleanup(func() {
		promptAttributionRecorder = previous
		metrics.SetDetailedTokenCollectionEnabled(false)
	})
	metrics.SetDetailedTokenCollectionEnabled(false)
	process := &Process{}
	records := []metrics.PromptCompileRecord{{
		AgentID: "agent-1", Runtime: "codex", LaunchID: "launch-1",
		SessionKind: "interactive", Freshness: "fresh",
		BlockID: "platform.agentteam.core", Durability: "persistent", Tokens: 42,
	}}

	process.SetLaunchPromptCompileMetrics(records)
	require.Empty(t, recorder.compile)
	require.False(t, process.contextExposureComplete)
	require.True(t, process.promptMetricsBaselineOK)

	metrics.SetDetailedTokenCollectionEnabled(true)
	process.InvalidateTokenAnalysisExposure()
	require.True(t, process.RecordRetainedPromptCompileMetrics())
	require.Equal(t, records, recorder.compile)
	require.False(t, process.RecordRetainedPromptCompileMetrics())
	require.Equal(t, records, recorder.compile)
}

func TestObservedDeliveryRestoresExposureCompletenessAfterLateCollectionEnable(t *testing.T) {
	ap := &Process{}
	metrics.SetDetailedTokenCollectionEnabled(false)
	ap.SetLaunchPromptCompileMetrics([]metrics.PromptCompileRecord{{
		BlockID: "platform.agentteam.core", Durability: "persistent", Tokens: 40,
	}})
	require.False(t, ap.contextExposureComplete)

	ap.activatePromptBoundary(promptMetricsBoundary{
		rendered: promptcompiler.RenderedPromptPayload{
			AttributionBlocks: []promptcompiler.PromptBlock{{
				ID: "event.current_message", Durability: promptcompiler.DurabilityTurnOnly, TokenEstimate: 10,
			}},
		},
	})

	require.True(t, ap.contextExposureComplete)
}

func TestObservedDeliveryDoesNotHideInvalidLaunchBaseline(t *testing.T) {
	ap := &Process{}
	metrics.SetDetailedTokenCollectionEnabled(true)
	t.Cleanup(func() { metrics.SetDetailedTokenCollectionEnabled(false) })
	ap.SetLaunchPromptCompileMetrics([]metrics.PromptCompileRecord{{
		BlockID: "platform.agentteam.core", Tokens: 40,
	}})
	require.False(t, ap.promptMetricsBaselineOK)

	ap.activatePromptBoundary(promptMetricsBoundary{
		rendered: promptcompiler.RenderedPromptPayload{
			AttributionBlocks: []promptcompiler.PromptBlock{{
				ID: "event.current_message", Durability: promptcompiler.DurabilityTurnOnly, TokenEstimate: 10,
			}},
		},
	})

	require.False(t, ap.contextExposureComplete)
}

func TestDeliveryPromptMetricsFollowAckedBoundary(t *testing.T) {
	recorder := useDeliveryPromptRecorder(t)
	driver := &correlatedOwnershipDriver{reconcileResult: RuntimeInputReconciliationPending}
	ap := newInboxRetryTestProcess(driver, time.Hour)
	ap.SessionID = "thread-1"
	ap.isIdle = true
	ap.deliveryOwnershipStore = &memoryDeliveryOwnershipStore{}
	setTokenTestProcessDimensions(ap)
	t.Cleanup(func() { stopInboxRetryTestProcess(ap) })
	message := types.InboxMessage{MessageID: "msg-1", Target: "#ops", Content: "diagnose this"}

	require.NoError(t, ap.deliverMessageWithContextSerialized(context.Background(), message, false))
	require.Empty(t, recorder.compile, "transport acceptance is not model-visible evidence")

	require.True(t, ap.handleRuntimeInputAcked(drivercontract.ParsedEvent{
		RuntimeInputID: "msg-1",
		SessionID:      "thread-1",
	}))
	require.NotEmpty(t, recorder.compile)
	count := len(recorder.compile)
	require.True(t, ap.handleRuntimeInputAcked(drivercontract.ParsedEvent{
		RuntimeInputID: "msg-1",
		SessionID:      "thread-1",
	}))
	require.Len(t, recorder.compile, count)
}

func TestDeliveryPromptMetricsDiscardAuthoritativeAbsent(t *testing.T) {
	recorder := useDeliveryPromptRecorder(t)
	driver := &correlatedOwnershipDriver{reconcileResult: RuntimeInputReconciliationAbsent}
	ap := newInboxRetryTestProcess(driver, time.Hour)
	ap.SessionID = "thread-1"
	ap.isIdle = true
	ap.deliveryOwnershipStore = &memoryDeliveryOwnershipStore{}
	setTokenTestProcessDimensions(ap)
	t.Cleanup(func() { stopInboxRetryTestProcess(ap) })

	require.NoError(t, ap.deliverMessageWithContextSerialized(context.Background(), types.InboxMessage{
		MessageID: "msg-absent", Target: "#ops", Content: "not acked",
	}, false))
	ap.mu.Lock()
	pending := ap.pendingRuntimeInputs["msg-absent"]
	ap.mu.Unlock()
	ap.runPendingRuntimeInputReconcile("msg-absent", pending.CreatedAt, 1, true)

	require.Empty(t, recorder.compile)
	ap.mu.Lock()
	_, staged := ap.pendingPromptMetrics["msg-absent"]
	ap.mu.Unlock()
	require.False(t, staged)
}

func TestDeliveryPromptMetricsRecordAmbiguousPositiveStdinWrite(t *testing.T) {
	recorder := useDeliveryPromptRecorder(t)
	stdin := &retiringShortWriteCloser{closed: make(chan struct{})}
	ap := newRawStdinOwnershipTestProcess(&memoryDeliveryOwnershipStore{}, stdin)
	setTokenTestProcessDimensions(ap)

	require.NoError(t, ap.DeliverForHandoff(context.Background(), types.InboxMessage{
		MessageID: "msg-partial", Target: "#ops", Content: "possibly visible",
	}))
	require.NotEmpty(t, recorder.compile)
}

func TestDeliveryPromptMetricsDiscardDefinitiveFailure(t *testing.T) {
	recorder := useDeliveryPromptRecorder(t)
	driver := &correlatedOwnershipDriver{
		runtimeErr:      errors.New("rejected"),
		reconcileResult: RuntimeInputReconciliationAbsent,
	}
	ap := newInboxRetryTestProcess(driver, time.Hour)
	ap.SessionID = "thread-1"
	ap.isIdle = true
	ap.deliveryOwnershipStore = &memoryDeliveryOwnershipStore{}
	setTokenTestProcessDimensions(ap)
	t.Cleanup(func() { stopInboxRetryTestProcess(ap) })

	require.NoError(t, ap.deliverMessageWithContextSerialized(context.Background(), types.InboxMessage{
		MessageID: "msg-rejected", Target: "#ops", Content: "reject this",
	}, false))
	require.Empty(t, recorder.compile)
}

func TestGeneratedPromptMetricsRecordOnlyAfterCollectionIsEnabled(t *testing.T) {
	recorder := useDeliveryPromptRecorder(t)
	ap := &Process{
		AgentID: "agent-1",
		processRuntimeBinding: processRuntimeBinding{
			driver: deliveryTestDriver{},
		},
	}
	setTokenTestProcessDimensions(ap)

	ap.recordGeneratedPromptCompileMetrics(promptcompiler.EventFinalVisibleReplyCheck, promptcompiler.GeneratedFinalReplyCheck)
	require.Len(t, recorder.compile, 1)
	require.Equal(t, "agentteam.final_visible_reply_check", recorder.compile[0].BlockID)
	require.Equal(t, string(promptcompiler.EventFinalVisibleReplyCheck), recorder.compile[0].CompileEvent)
	require.Positive(t, recorder.compile[0].Tokens)
	encoded, err := json.Marshal(recorder.compile[0])
	require.NoError(t, err)
	require.NotContains(t, string(encoded), "Send the pending visible reply.")
	require.Equal(t, map[string]metrics.InvocationExposureBlock{
		"agentteam.final_visible_reply_check\x00turn_only": {
			BlockID: "agentteam.final_visible_reply_check", Durability: "turn_only",
			Tokens: recorder.compile[0].Tokens,
		},
	}, ap.activePromptMetrics)

	metrics.SetDetailedTokenCollectionEnabled(false)
	ap.recordGeneratedPromptCompileMetrics(promptcompiler.EventFinalVisibleReplyCheck, promptcompiler.GeneratedFinalReplyCheck)
	require.Len(t, recorder.compile, 1)
}

func TestDeliveryAttachmentMetricJoinsAckedBoundaryWithoutBody(t *testing.T) {
	recorder := useDeliveryPromptRecorder(t)
	ap := &Process{
		AgentID: "agent-1",
		processRuntimeBinding: processRuntimeBinding{
			driver: deliveryTestDriver{},
		},
	}
	setTokenTestProcessDimensions(ap)
	message := types.InboxMessage{
		MessageID:      "msg-attachment",
		Content:        "inspect this",
		AttachmentList: []types.MessageAttachment{{ID: "att-1", Filename: "secret.csv"}},
	}
	_, err := ap.renderDeliveryTextWithContext(context.Background(), message, "stdin_delivery", "")
	require.NoError(t, err)
	ap.SetRuntimeAttachmentResolver(testRuntimeAttachmentResolver{
		attachments: []types.RuntimeAttachment{{
			ID: "att-1", Filename: "secret.csv", Path: "/tmp/private/secret.csv",
		}},
	})
	compiled, err := ap.compileDeliveryPromptWithContext(context.Background(), message, "turn_direct", "")
	require.NoError(t, err)
	input := ap.runtimeInputForDelivery(context.Background(), message, compiled)
	require.Contains(t, input.PromptText(), "/tmp/private/secret.csv")
	ap.commitDeliveryPromptCompileMetrics(message.MessageID)

	require.Len(t, recorder.compile, 3)
	var attachment metrics.PromptCompileRecord
	for _, record := range recorder.compile {
		if record.BlockID == "runtime_attachment_manifest" {
			attachment = record
		}
	}
	require.Equal(t, "runtime_attachment_manifest", attachment.BlockID)
	require.Positive(t, attachment.Tokens)
	encoded, err := json.Marshal(recorder.compile)
	require.NoError(t, err)
	require.NotContains(t, string(encoded), "/tmp/private/secret.csv")
	require.NotContains(t, string(encoded), "inspect this")
}

func TestAttachmentThenNormalPayloadAttributionKeepsUnavailableState(t *testing.T) {
	recorder := useDeliveryPromptRecorder(t)
	SetTaskPayloadFingerprinter(func(payload []byte) string {
		return "keyed:" + string(payload)
	})
	t.Cleanup(func() { SetTaskPayloadFingerprinter(nil) })
	ap := &Process{
		AgentID: "agent-1",
		processRuntimeBinding: processRuntimeBinding{
			driver: deliveryTestDriver{},
		},
	}
	setTokenTestProcessDimensions(ap)
	compiler := promptcompiler.New(promptcompiler.DefaultPolicy())

	compileAndRecord := func(messageID, content string, attachments []types.MessageAttachment) {
		t.Helper()
		compiled, err := compiler.CompileDeliveryEnvelope(t.Context(), promptcompiler.EventRoutedMessageDelivery, messageID, promptcompiler.CompileInput{
			AgentID: "agent-1",
			Runtime: "claude",
			Delivery: promptcompiler.DeliveryCompileInput{
				MessageID:      messageID,
				Target:         "#ops",
				ChannelType:    "channel",
				AuthorIdentity: types.AuthorIdentity{ActorType: "human", DisplayName: "alice"},
				Content:        content,
				MessageType:    "text",
				Attachments:    attachments,
			},
		})
		require.NoError(t, err)
		rendered, err := compiler.RenderRuntimeInput(t.Context(), compiled, "claude", promptcompiler.RenderOptions{})
		require.NoError(t, err)
		ap.recordDeliveryPromptCompileMetrics(compiled, rendered)
	}

	compileAndRecord("with-attachment", "Inspect this.", []types.MessageAttachment{{
		ID: "att-1", Filename: "secret.csv",
	}})
	compileAndRecord("without-attachment", "Inspect this.", nil)

	var unavailable, identified bool
	for _, record := range recorder.compile {
		if !record.TaskPayloadPresent {
			continue
		}
		if record.TaskPayloadFP == "" {
			unavailable = true
		} else {
			identified = true
		}
	}
	require.True(t, unavailable)
	require.True(t, identified)
}

func useDeliveryPromptRecorder(t *testing.T) *promptAttributionRecorderStub {
	t.Helper()
	recorder := &promptAttributionRecorderStub{}
	previous := promptAttributionRecorder
	promptAttributionRecorder = func() metrics.PromptAttributionRecorder { return recorder }
	metrics.SetDetailedTokenCollectionEnabled(true)
	t.Cleanup(func() {
		promptAttributionRecorder = previous
		metrics.SetDetailedTokenCollectionEnabled(false)
	})
	return recorder
}

func setTokenTestProcessDimensions(ap *Process) {
	ap.processUsageState = processUsageState{
		launchID:         "launch-test",
		runtime:          "claude",
		sessionKind:      "interactive",
		sessionFreshness: metrics.PromptSessionFreshnessFresh,
	}
}

func TestSessionFreshnessForLaunch(t *testing.T) {
	require.Equal(t, "fresh", sessionFreshnessForLaunch(drivercontract.LaunchConfig{}))
	require.Equal(t, "resume", sessionFreshnessForLaunch(drivercontract.LaunchConfig{SessionID: "thread-1"}))
}

func TestSessionResumeFallbackDetailIsStructuredAndConsistent(t *testing.T) {
	for _, runtime := range []runtimeconfig.RuntimeType{
		runtimeconfig.TypeCodex,
		runtimeconfig.TypeOpenCode,
		runtimeconfig.TypeKimi,
		runtimeconfig.TypePi,
	} {
		t.Run(string(runtime), func(t *testing.T) {
			detail := sessionResumeFallbackDetail(
				drivercontract.LaunchConfig{RuntimeType: runtime},
				drivercontract.SessionResumeFallback{Reason: "invalid_or_not_found_session", Operation: "session/preflight"},
			)
			require.JSONEq(t, `{"attempt":1,"code":"provider_session_fresh_fallback","operation":"session/preflight","reason":"invalid_or_not_found_session","runtime":"`+string(runtime)+`"}`, detail)
		})
	}
}