package metrics

import (
	"context"
	"fmt"
	"path/filepath"
	"testing"
	"time"

	"github.com/prometheus/prometheus/model/labels"
	"github.com/prometheus/prometheus/tsdb"
	"github.com/prometheus/prometheus/tsdb/chunkenc"
	"github.com/stretchr/testify/require"
)

func newTestTokenTrackingStore(t *testing.T) *TokenTrackingStore {
	t.Helper()
	path := filepath.Join(t.TempDir(), tokenTrackingFilename)
	store, err := openTokenTrackingStore(path, true)
	require.NoError(t, err)
	t.Cleanup(func() { require.NoError(t, store.Close()) })
	return store
}

func TestTokenTrackingStoreUsesIndependentSevenDayDatabase(t *testing.T) {
	require.Equal(t, "tokentracking.tsdb", filepath.Base(tokenTrackingPath))
	require.NotEqual(t, filepath.Clean(tsdbPath), filepath.Clean(tokenTrackingPath))
	require.Equal(t, 7*24*time.Hour, tokenTrackingRetention)
	require.Equal(t, 2*time.Hour, tokenTrackingBlock)

	store := newTestTokenTrackingStore(t)
	require.Equal(t, tokenTrackingFilename, filepath.Base(store.path))
	require.NotSame(t, store.db, globalTSDB)
}

func TestTokenTrackingStoreWritesAlignedSanitizedSamples(t *testing.T) {
	store := newTestTokenTrackingStore(t)
	timestamp := time.Now().Add(-time.Minute).UnixMilli()
	store.RecordTokenTracking(TokenSourceObservation{
		Timestamp:        timestamp,
		AgentID:          "agent-a",
		Runtime:          "Codex",
		Model:            "GPT-5.2",
		SourceKind:       TokenSourceKindTool,
		SourceName:       "agentteam.memory_search",
		InputToken:       100,
		CachedInputToken: 40,
		OutputToken:      20,
	})
	// The store serializes duplicate timestamps across calls instead of letting
	// Prometheus silently replace the second observation.
	store.RecordTokenTracking(TokenSourceObservation{
		Timestamp:        timestamp,
		AgentID:          "agent-a",
		Runtime:          "codex",
		Model:            "gpt-5.2",
		SourceKind:       TokenSourceKindCommand,
		SourceName:       "git",
		InputToken:       10,
		CachedInputToken: 4,
		OutputToken:      2,
	})

	result := store.QueryTokenTracking(TokenSourceQuery{
		AgentIDs: []string{"agent-a"},
		Window:   TimeWindow6Hours,
	})
	require.True(t, result.Available)
	require.Len(t, result.Trend.Series, 2)

	for _, metricName := range []string{
		metricTokenTrackingCalls,
		metricTokenTrackingInput,
		metricTokenTrackingCached,
		metricTokenTrackingOutput,
	} {
		samples := tokenTrackingMetricSamples(t, store.db, metricName)
		require.Len(t, samples, 2)
		for _, sample := range samples {
			names := make([]string, 0, sample.Labels().Len())
			sample.Labels().Range(func(label labels.Label) {
				names = append(names, label.Name)
			})
			require.ElementsMatch(t, []string{
				"__name__",
				labelAgentID,
				labelTokenSourceKind,
				labelTokenSourceName,
			}, names)
		}
	}
}

func TestTokenTrackingStoreRejectsRawOrUnboundedSourceLabels(t *testing.T) {
	store := newTestTokenTrackingStore(t)
	store.RecordTokenTracking(TokenSourceObservation{
		AgentID:     "agent-a",
		Runtime:     "codex",
		Model:       "gpt-5",
		SourceKind:  TokenSourceKindCommand,
		SourceName:  "grep -n secret /private/path",
		InputToken:  10,
		OutputToken: 2,
	})

	result := store.QueryTokenTracking(TokenSourceQuery{AgentIDs: []string{"agent-a"}, Window: TimeWindow1Day})
	require.Len(t, result.Trend.Series, 1)
	require.Equal(t, TokenSourceKindUnknown, result.Trend.Series[0].SourceKind)
	require.Equal(t, TokenSourceKindUnknown, result.Trend.Series[0].Label)
	require.NotContains(t, fmt.Sprint(result), "secret")
	require.NotContains(t, fmt.Sprint(result), "/private/path")
}

func TestTokenTrackingQueryAggregatesSourcesAndTrend(t *testing.T) {
	store := newTestTokenTrackingStore(t)
	now := time.Now().Add(-time.Minute)
	for index, observation := range []TokenSourceObservation{
		{AgentID: "agent-a", Runtime: "codex", Model: "gpt-5", SourceKind: TokenSourceKindCommand, SourceName: "grep", InputToken: 100, CachedInputToken: 60, OutputToken: 20},
		{AgentID: "agent-a", Runtime: "codex", Model: "gpt-5", SourceKind: TokenSourceKindTool, SourceName: "agentteam.message_send", InputToken: 500, CachedInputToken: 300, OutputToken: 100},
	} {
		observation.Timestamp = now.Add(time.Duration(index) * time.Millisecond).UnixMilli()
		store.RecordTokenTracking(observation)
	}

	result := store.QueryTokenTracking(TokenSourceQuery{AgentIDs: []string{"agent-a"}, Window: TimeWindow1Day})
	require.True(t, result.Available)
	// Series order is the server-owned total-token ranking.
	require.Len(t, result.Trend.Series, 2)
	require.Equal(t, []string{"agentteam.message_send", "grep"}, []string{
		result.Trend.Series[0].Label,
		result.Trend.Series[1].Label,
	})
	for _, series := range result.Trend.Series {
		require.NotEmpty(t, series.Key)
		require.Contains(t, series.Key, tokenSourceTrendKeyPrefix)
	}
	require.NotEmpty(t, result.Trend.Data)
	var previousTimestamp int64
	for _, point := range result.Trend.Data {
		require.NotZero(t, point.Timestamp, "trend points carry their bucket timestamp")
		require.Greater(t, point.Timestamp, previousTimestamp, "trend points are ordered by bucket")
		previousTimestamp = point.Timestamp
	}
}

// TestTokenTrackingQueryTrendDistinguishesSameNameDifferentKind verifies the
// trend aggregates by (sourceKind, sourceName), so a tool and a command that
// share a name stay separate series instead of merging.
func TestTokenTrackingQueryTrendDistinguishesSameNameDifferentKind(t *testing.T) {
	store := newTestTokenTrackingStore(t)
	now := time.Now().Add(-time.Minute).UnixMilli()
	store.RecordTokenTracking(TokenSourceObservation{
		Timestamp: now, AgentID: "agent-a", Runtime: "codex", Model: "gpt-5",
		SourceKind: TokenSourceKindTool, SourceName: "shared", InputToken: 40, OutputToken: 10,
	})
	store.RecordTokenTracking(TokenSourceObservation{
		Timestamp: now + 1, AgentID: "agent-a", Runtime: "codex", Model: "gpt-5",
		SourceKind: TokenSourceKindCommand, SourceName: "shared", InputToken: 20, OutputToken: 5,
	})
	result := store.QueryTokenTracking(TokenSourceQuery{AgentIDs: []string{"agent-a"}, Window: TimeWindow1Day})
	require.Len(t, result.Trend.Series, 2)
	kinds := make([]string, 0, 2)
	for _, series := range result.Trend.Series {
		require.Equal(t, "shared", series.Label)
		kinds = append(kinds, series.SourceKind)
	}
	require.ElementsMatch(t, []string{TokenSourceKindTool, TokenSourceKindCommand}, kinds)
}

// TestTokenTrackingQueryTrendUsesOneGlobalRemainder verifies that all sources
// outside the global Top 15 share one bounded series without losing volume.
func TestTokenTrackingQueryTrendUsesOneGlobalRemainder(t *testing.T) {
	store := newTestTokenTrackingStore(t)
	now := time.Now().Add(-time.Minute).UnixMilli()
	var expectedTotal uint64
	// Fill Top N tools so direct + unknown land in the remainder.
	for index := 0; index < tokenSourceTrendTopN; index++ {
		total := uint64(100-index) + 1
		expectedTotal += total
		store.RecordTokenTracking(TokenSourceObservation{
			Timestamp: now + int64(index), AgentID: "agent-a", Runtime: "codex", Model: "gpt-5",
			SourceKind: TokenSourceKindTool, SourceName: fmt.Sprintf("tool_%02d", index),
			InputToken: uint64(100 - index), OutputToken: 1,
		})
	}
	store.RecordTokenTracking(TokenSourceObservation{
		Timestamp: now + 100, AgentID: "agent-a", Runtime: "codex", Model: "gpt-5",
		SourceKind: TokenSourceKindDirect, SourceName: "direct_a", InputToken: 30, OutputToken: 5,
	})
	expectedTotal += 35
	store.RecordTokenTracking(TokenSourceObservation{
		Timestamp: now + 101, AgentID: "agent-a", Runtime: "codex", Model: "gpt-5",
		SourceKind: TokenSourceKindUnknown, SourceName: TokenSourceKindUnknown, InputToken: 20, OutputToken: 5,
	})
	expectedTotal += 25
	result := store.QueryTokenTracking(TokenSourceQuery{AgentIDs: []string{"agent-a"}, Window: TimeWindow1Day})
	require.Len(t, result.Trend.Series, tokenSourceTrendTopN+1)
	require.Equal(t, tokenSourceTrendRemainderKey, result.Trend.Series[tokenSourceTrendTopN].Key)
	require.Empty(t, result.Trend.Series[tokenSourceTrendTopN].SourceKind)
	// Total across the trend window is preserved (no value lost to collisions).
	var trendTotal uint64
	for _, point := range result.Trend.Data {
		for _, value := range point.Values {
			trendTotal += value
		}
	}
	require.Equal(t, expectedTotal, trendTotal)
}

func TestTokenTrackingStoreAcceptsContextRemainderSource(t *testing.T) {
	store := newTestTokenTrackingStore(t)
	store.RecordTokenTracking(TokenSourceObservation{
		AgentID:          "agent-a",
		Runtime:          "codex",
		Model:            "gpt-5",
		SourceKind:       TokenSourceKindContext,
		SourceName:       TokenSourceKindContext,
		InputToken:       30,
		CachedInputToken: 20,
		OutputToken:      10,
	})

	result := store.QueryTokenTracking(TokenSourceQuery{
		AgentIDs: []string{"agent-a"},
		Window:   TimeWindow1Day,
	})
	require.Len(t, result.Trend.Series, 1)
	require.Equal(t, TokenSourceKindContext, result.Trend.Series[0].SourceKind)
	require.Equal(t, TokenSourceKindContext, result.Trend.Series[0].Label)
}

func TestTokenTrackingQueryRollsUpTrendBeyondTopN(t *testing.T) {
	store := newTestTokenTrackingStore(t)
	now := time.Now().Add(-time.Minute)
	for index := 0; index < tokenSourceTrendTopN+3; index++ {
		store.RecordTokenTracking(TokenSourceObservation{
			Timestamp:   now.Add(time.Duration(index) * time.Millisecond).UnixMilli(),
			AgentID:     "agent-a",
			Runtime:     "codex",
			Model:       "gpt-5",
			SourceKind:  TokenSourceKindTool,
			SourceName:  fmt.Sprintf("tool_%02d", index),
			InputToken:  uint64(tokenSourceTrendTopN + 3 - index),
			OutputToken: 1,
		})
	}
	result := store.QueryTokenTracking(TokenSourceQuery{AgentIDs: []string{"agent-a"}, Window: TimeWindow1Day})
	// Trend series keeps Top 15 named sources plus one global roll-up.
	require.Len(t, result.Trend.Series, tokenSourceTrendTopN+1)
	require.Equal(t, tokenSourceTrendRemainderKey, result.Trend.Series[len(result.Trend.Series)-1].Key)
}

func TestTokenTrackingQueryUnavailableAndEmptyCollections(t *testing.T) {
	result := (*TokenTrackingStore)(nil).QueryTokenTracking(TokenSourceQuery{Window: TimeWindow1Day})
	require.False(t, result.Available)
	require.NotNil(t, result.Trend.Data)

	store := newTestTokenTrackingStore(t)
	empty := store.QueryTokenTracking(TokenSourceQuery{AgentIDs: nil, Window: TimeWindow1Day})
	require.True(t, empty.Available)
	require.NotNil(t, empty.Trend.Data)
}

// TestTokenTrackingQueryRejectsSamplesWithoutLabels verifies the
// query aggregates only series carrying a resolved source_kind label. A
// series with a non-resolved kind is rejected so it can neither pollute
// totals nor the per-source trend.
func TestTokenTrackingQueryRejectsSamplesWithoutLabels(t *testing.T) {
	store := newTestTokenTrackingStore(t)
	now := time.Now().Add(-time.Minute).UnixMilli()

	appender := store.db.Appender(context.Background())
	for _, metricName := range []string{
		metricTokenTrackingCalls,
		metricTokenTrackingInput,
		metricTokenTrackingCached,
		metricTokenTrackingOutput,
	} {
		invalidKindLabels := labels.FromStrings(
			"__name__", metricName,
			labelAgentID, "agent-a",
			labelTokenSourceKind, "invalid",
			labelTokenSourceName, "invalid_source",
		)
		_, err := appender.Append(0, invalidKindLabels, now, 10)
		require.NoError(t, err)
		// Valid kind but path-like source_name: the read path must also reject it.
		pathNameLabels := labels.FromStrings(
			"__name__", metricName,
			labelAgentID, "agent-a",
			labelTokenSourceKind, TokenSourceKindTool,
			labelTokenSourceName, "/private/path",
		)
		_, err = appender.Append(0, pathNameLabels, now, 10)
		require.NoError(t, err)
	}
	require.NoError(t, appender.Commit())

	store.RecordTokenTracking(TokenSourceObservation{
		Timestamp:   now + 1,
		AgentID:     "agent-a",
		Runtime:     "codex",
		Model:       "gpt-5",
		SourceKind:  TokenSourceKindCommand,
		SourceName:  "grep",
		InputToken:  100,
		OutputToken: 20,
	})

	result := store.QueryTokenTracking(TokenSourceQuery{AgentIDs: []string{"agent-a"}, Window: TimeWindow1Day})
	require.True(t, result.Available)
	require.Len(t, result.Trend.Series, 1, "only resolved-label series may aggregate")
	require.Equal(t, "grep", result.Trend.Series[0].Label)
}

// TestTokenTrackingStoreRejectsPathLikeSourceName verifies the persistence
// layer fails closed on path-like source names: normalizedTokenSourceName
// rejects '/', so a source name that is a bare path is collapsed to unknown
// and never written to the TSDB or returned to the UI.
func TestTokenTrackingStoreRejectsPathLikeSourceName(t *testing.T) {
	store := newTestTokenTrackingStore(t)
	store.RecordTokenTracking(TokenSourceObservation{
		AgentID:     "agent-a",
		Runtime:     "codex",
		Model:       "gpt-5",
		SourceKind:  TokenSourceKindTool,
		SourceName:  "/private/path",
		InputToken:  10,
		OutputToken: 2,
	})
	result := store.QueryTokenTracking(TokenSourceQuery{AgentIDs: []string{"agent-a"}, Window: TimeWindow1Day})
	require.Len(t, result.Trend.Series, 1)
	require.Equal(t, TokenSourceKindUnknown, result.Trend.Series[0].SourceKind)
	require.NotContains(t, fmt.Sprint(result), "/private/path")
}

func TestTokenTrackingQueryTrendSeriesIsBoundedToTop15AndOther(t *testing.T) {
	store := newTestTokenTrackingStore(t)
	now := time.Now().Add(-time.Minute).UnixMilli()
	kinds := []string{TokenSourceKindTool, TokenSourceKindCommand, TokenSourceKindContext, TokenSourceKindDirect, TokenSourceKindMultiple, TokenSourceKindUnknown}
	for index := 0; index < tokenSourceTrendTopN+2; index++ {
		for _, kind := range kinds {
			store.RecordTokenTracking(TokenSourceObservation{
				Timestamp:   now + int64(index),
				AgentID:     "agent-a",
				Runtime:     "codex",
				Model:       "gpt-5",
				SourceKind:  kind,
				SourceName:  fmt.Sprintf("%s_%02d", kind, index),
				InputToken:  uint64(tokenSourceTrendTopN + 2 - index),
				OutputToken: 1,
			})
		}
	}
	result := store.QueryTokenTracking(TokenSourceQuery{AgentIDs: []string{"agent-a"}, Window: TimeWindow1Day})
	require.Len(t, result.Trend.Series, tokenSourceTrendTopN+1)
	require.Equal(t, tokenSourceTrendRemainderKey, result.Trend.Series[tokenSourceTrendTopN].Key)
}

type tokenTrackingSeriesSample struct {
	labels labels.Labels
}

func (sample tokenTrackingSeriesSample) Labels() labels.Labels { return sample.labels }

func tokenTrackingMetricSamples(t *testing.T, db *tsdb.DB, metricName string) []tokenTrackingSeriesSample {
	t.Helper()
	querier, err := db.Querier(time.Now().Add(-time.Hour).UnixMilli(), time.Now().Add(time.Hour).UnixMilli())
	require.NoError(t, err)
	defer querier.Close()
	set := querier.Select(
		context.Background(),
		false,
		nil,
		labels.MustNewMatcher(labels.MatchEqual, "__name__", metricName),
	)
	var out []tokenTrackingSeriesSample
	for set.Next() {
		series := set.At()
		iterator := series.Iterator(nil)
		for iterator.Next() == chunkenc.ValFloat {
			out = append(out, tokenTrackingSeriesSample{labels: series.Labels()})
		}
		require.NoError(t, iterator.Err())
	}
	require.NoError(t, set.Err())
	return out
}