package metrics

import (
	"testing"
	"time"

	"github.com/stretchr/testify/require"
)

func TestPromptCompileMetricsRoundTrip(t *testing.T) {
	tsdbCol, _ := newTestTSDB(t)
	defer tsdbCol.Close()
	collector := tsdbCol
	collector.RecordPromptCompileMetrics([]PromptCompileRecord{{
		Timestamp:          time.Now().UnixMilli(),
		AgentID:            "agent-1",
		Runtime:            "codex",
		LaunchID:           "launch-1",
		SessionKind:        "interactive",
		Freshness:          PromptSessionFreshnessFresh,
		CompileEvent:       "thread_start",
		Stage:              "thread_contract",
		Slot:               "baseInstructions",
		BlockID:            "platform.agentteam.startup",
		BlockLayer:         "platform",
		BlockSlot:          "base",
		BlockType:          "instruction",
		TaskPayloadPresent: true,
		TaskPayloadFP:      "payload-fingerprint",
		Tokens:             123,
	}})
	records, _, _ := tsdbCol.queryPromptCompileRecordsBounded(
		promptCompileQuery{Window: TimeWindow6Hours},
		10,
	)
	require.NotEmpty(t, records)
	last := records[len(records)-1]
	require.Equal(t, "launch-1", last.LaunchID)
	require.Equal(t, "thread_contract", last.Stage)
	require.True(t, last.TaskPayloadPresent)
	require.Equal(t, "payload-fingerprint", last.TaskPayloadFP)
	require.Equal(t, uint64(1), last.Records)
	require.Equal(t, uint64(123), last.Tokens)
}

func TestPromptCompileMetricsRoundTripPreservesUnavailableTaskPayloadIdentity(t *testing.T) {
	tsdbCol, _ := newTestTSDB(t)
	defer tsdbCol.Close()
	tsdbCol.RecordPromptCompileMetrics([]PromptCompileRecord{{
		Timestamp:          time.Now().UnixMilli(),
		AgentID:            "agent-1",
		Runtime:            "codex",
		LaunchID:           "launch-unavailable",
		SessionKind:        "interactive",
		Freshness:          PromptSessionFreshnessFresh,
		CompileEvent:       "wake",
		Stage:              "turn_context",
		Slot:               "turn",
		BlockID:            "event.current_message",
		TaskPayloadPresent: true,
		Tokens:             10,
	}})

	records, _, _ := tsdbCol.queryPromptCompileRecordsBounded(
		promptCompileQuery{Window: TimeWindow6Hours, LaunchID: "launch-unavailable"},
		100,
	)
	require.Len(t, records, 1)
	require.True(t, records[0].TaskPayloadPresent)
	require.Empty(t, records[0].TaskPayloadFP)
}

func TestTokenAnalysisEvidenceRoundTripAndBound(t *testing.T) {
	tsdbCol, _ := newTestTSDB(t)
	defer tsdbCol.Close()
	now := time.Now().Add(-time.Minute).UnixMilli()
	startedAt := now - time.Minute.Milliseconds()
	tsdbCol.RecordPromptCompileMetrics([]PromptCompileRecord{{
		Timestamp:       now,
		LaunchStartedAt: startedAt,
		AgentID:         "agent-1",
		Runtime:         "codex",
		LaunchID:        "launch-1",
		SessionKind:     "interactive",
		Freshness:       PromptSessionFreshnessFresh,
		EventID:         "event-1",
		CompileEvent:    "wake",
		Stage:           "delivery_envelope",
		Slot:            "turn",
		BlockID:         "event.wake_context",
		BlockLayer:      "event",
		BlockSlot:       "turn",
		BlockType:       "turn_context",
		Durability:      "turn_only",
		Tokens:          123,
	}})
	for index := 0; index < 2; index++ {
		exposure := InvocationExposureRecord{
			Complete: true,
			Blocks: []InvocationExposureBlock{{
				BlockID: "platform.agentteam.core", Durability: "persistent", Tokens: 42,
			}},
		}
		tsdbCol.RecordUsageObservation(UsageObservationRecord{
			Timestamp:       now + int64(index+1),
			LaunchStartedAt: startedAt,
			ObservationID:   "observation-" + string(rune('a'+index)),
			AgentID:         "agent-1",
			Runtime:         "codex",
			LaunchID:        "launch-1",
			SessionKind:     "interactive",
			SessionID:       "thread-1",
			Model:           "gpt",
			InvocationID:    "turn-" + string(rune('a'+index)),
			InvocationCause: "unknown",
			InputTokens:     uint64(100 + index),
			CachedTokens:    uint64(10 + index),
			OutputTokens:    uint64(20 + index),
		}, &exposure)
	}

	evidence := tsdbCol.QueryTokenAnalysisEvidence(TokenEvidenceQuery{
		Window:   TimeWindow6Hours,
		AgentIDs: []string{"agent-1"},
		LaunchID: "launch-1",
		Limit:    10,
	})
	require.Len(t, evidence.PromptRecords, 1)
	require.Equal(t, "event-1", evidence.PromptRecords[0].EventID)
	require.Equal(t, startedAt, evidence.PromptRecords[0].LaunchStartedAt)
	require.Equal(t, "turn_only", evidence.PromptRecords[0].Durability)
	require.Len(t, evidence.Usage, 2)
	require.Equal(t, uint64(100), evidence.Usage[0].InputTokens)
	require.Equal(t, startedAt, evidence.Usage[0].LaunchStartedAt)
	require.Equal(t, uint64(10), evidence.Usage[0].CachedTokens)
	require.Equal(t, uint64(20), evidence.Usage[0].OutputTokens)
	require.Equal(t, []ExposureCoverageRecord{{
		LaunchID: "launch-1", Observations: 2, CompleteObservations: 2,
	}}, evidence.ExposureCoverage)
	require.Equal(t, []BlockExposureRecord{{
		LaunchID: "launch-1", BlockID: "platform.agentteam.core", Durability: "persistent",
		ExposureCount: 2, Tokens: 84,
	}}, evidence.BlockExposures)
	require.LessOrEqual(t, evidence.Scanned, evidence.Limit)

	bounded := tsdbCol.QueryTokenAnalysisEvidence(TokenEvidenceQuery{
		Window:   TimeWindow6Hours,
		AgentIDs: []string{"agent-1"},
		Limit:    2,
	})
	require.True(t, bounded.Truncated)
	require.LessOrEqual(t, bounded.Scanned, bounded.Limit)
}

func TestInvocationExposureRoundTripDeduplicatesRuntimeReplay(t *testing.T) {
	tsdbCol, _ := newTestTSDB(t)
	defer tsdbCol.Close()
	now := time.Now().Add(-time.Minute).UnixMilli()
	item := InvocationExposureRecord{
		Complete: true,
		Blocks: []InvocationExposureBlock{{
			BlockID: "platform.agentteam.core", Durability: "persistent", Tokens: 42,
		}},
	}
	tsdbCol.RecordUsageObservation(UsageObservationRecord{
		Timestamp:       now,
		AgentID:         "agent-1",
		Runtime:         "opencode",
		LaunchID:        "launch-1",
		SessionKind:     "interactive",
		SessionID:       "session-1",
		InvocationID:    "step-1",
		InvocationCause: "unknown",
		InputTokens:     100,
	}, &item)
	tsdbCol.RecordUsageObservation(UsageObservationRecord{
		Timestamp:       now + 1,
		AgentID:         "agent-1",
		Runtime:         "opencode",
		LaunchID:        "launch-1",
		SessionKind:     "interactive",
		SessionID:       "session-1",
		InvocationID:    "step-1",
		InvocationCause: "unknown",
		InputTokens:     100,
	}, &item)
	tsdbCol.RecordUsageObservation(UsageObservationRecord{
		Timestamp:        now + 2,
		AgentID:          "agent-1",
		Runtime:          "kimi",
		LaunchID:         "launch-1",
		SessionKind:      "interactive",
		SessionID:        "session-1",
		InvocationID:     "subagent-step-1",
		InvocationCause:  "unknown",
		ExposureExcluded: true,
		InputTokens:      50,
	}, &item)

	evidence := tsdbCol.QueryTokenAnalysisEvidence(TokenEvidenceQuery{
		Window:   TimeWindow6Hours,
		AgentIDs: []string{"agent-1"},
		Limit:    100,
	})
	require.Equal(t, []ExposureCoverageRecord{{
		LaunchID: "launch-1", Observations: 1, CompleteObservations: 1,
	}}, evidence.ExposureCoverage)
	require.Equal(t, []BlockExposureRecord{{
		LaunchID: "launch-1", BlockID: "platform.agentteam.core", Durability: "persistent",
		ExposureCount: 1, Tokens: 42,
	}}, evidence.BlockExposures)
	require.Len(t, evidence.Usage, 3)
	require.True(t, evidence.Usage[2].ExposureExcluded)
}