package metrics

import (
	"context"
	"fmt"
	"log/slog"
	"regexp"
	"sort"
	"strings"
	"time"

	"github.com/prometheus/prometheus/model/labels"
	"github.com/prometheus/prometheus/storage"
	"github.com/prometheus/prometheus/tsdb/chunkenc"
)

const (
	maxTokenTrackingQuerySamples = 250_000
	tokenSourceTrendRemainderKey = "other"

	// tokenSourceTrendKeyPrefix is the opaque prefix for Top N series keys. It
	// is never the raw source name, so a real source named "other" cannot
	// collide with a roll-up key and identical names of different kinds stay
	// distinct.
	tokenSourceTrendKeyPrefix = "source_"
)

type tokenTrackingSampleKey struct {
	timestamp int64
	agentID   string
	kind      string
	name      string
}

type tokenTrackingSample struct {
	key      tokenTrackingSampleKey
	presence uint8
	input    uint64
	output   uint64
}

// tokenSourceIdentity is a (kind, name) trend identity. Trend aggregates by
// identity so the same name of different kinds is never merged.
type tokenSourceIdentity struct {
	kind string
	name string
}

// QueryTokenTracking returns a bounded aggregate page for the explicit Agent set.
func (store *TokenTrackingStore) QueryTokenTracking(query TokenSourceQuery) TokenSourceResult {
	if store == nil {
		return emptyTokenSourceResult(query.Window, false)
	}
	selectedAgentIDs := normalizedTokenTrackingAgentIDs(query.AgentIDs)
	if len(selectedAgentIDs) == 0 {
		return emptyTokenSourceResult(query.Window, true)
	}

	store.mu.RLock()
	defer store.mu.RUnlock()
	if store.db == nil {
		return emptyTokenSourceResult(query.Window, false)
	}

	now := time.Now()
	querier, err := store.db.Querier(now.Add(-query.Window.Duration()).UnixMilli(), now.UnixMilli())
	if err != nil {
		slog.Warn("failed to query token tracking TSDB", "error", err)
		return emptyTokenSourceResult(query.Window, false)
	}
	defer querier.Close()

	records := queryTokenTrackingRecords(querier, selectedAgentIDs)
	return aggregateTokenSourceResult(query.Window, records, now)
}

func queryTokenTrackingRecords(
	querier storage.Querier,
	selectedAgentIDs []string,
) map[tokenTrackingSampleKey]*tokenTrackingSample {
	records := make(map[tokenTrackingSampleKey]*tokenTrackingSample)
	agentMatcher := tokenTrackingValuesMatcher(labelAgentID, selectedAgentIDs)
	for metricIndex, metric := range []struct {
		name  string
		apply func(*tokenTrackingSample, uint64)
	}{
		{name: metricTokenTrackingInput, apply: func(sample *tokenTrackingSample, value uint64) { sample.input += value }},
		{name: metricTokenTrackingOutput, apply: func(sample *tokenTrackingSample, value uint64) { sample.output += value }},
	} {
		matchers := []*labels.Matcher{
			labels.MustNewMatcher(labels.MatchEqual, "__name__", metric.name),
			agentMatcher,
		}
		seriesSet := querier.Select(context.Background(), false, nil, matchers...)
		scanned := 0
		for seriesSet.Next() {
			series := seriesSet.At()
			lbls := series.Labels()
			// Aggregate only series carrying canonical source_kind and
			// source_name labels. The source_name check fail-closes the read
			// path: even a pre-existing or injected path-like name in the TSDB
			// is skipped instead of being returned to the UI. Anything else is
			// skipped before its iterator is read, so the sample cap below
			// counts only admitted samples and malformed series can neither
			// pollute totals nor the per-source trend.
			kind := lbls.Get(labelTokenSourceKind)
			if !validTokenSourceKind(kind) {
				continue
			}
			name := lbls.Get(labelTokenSourceName)
			if normalizedTokenSourceName(name, 120) != name {
				continue
			}
			iterator := series.Iterator(nil)
			for iterator.Next() == chunkenc.ValFloat {
				if scanned >= maxTokenTrackingQuerySamples {
					break
				}
				scanned++
				timestamp, value := iterator.At()
				key := tokenTrackingSampleKey{
					timestamp: timestamp,
					agentID:   lbls.Get(labelAgentID),
					kind:      kind,
					name:      name,
				}
				sample := records[key]
				if sample == nil {
					sample = &tokenTrackingSample{key: key}
					records[key] = sample
				}
				metric.apply(sample, uint64(value))
				sample.presence |= 1 << metricIndex
			}
			if scanned >= maxTokenTrackingQuerySamples {
				break
			}
		}
	}
	return records
}

func aggregateTokenSourceResult(
	window TimeWindow,
	records map[tokenTrackingSampleKey]*tokenTrackingSample,
	now time.Time,
) TokenSourceResult {
	interval := MachineDashboardAggregationInterval(window)
	result := emptyTokenSourceResult(window, true)

	bucketByIdentity := make(map[int64]map[tokenSourceIdentity]uint64)
	identityTotal := make(map[tokenSourceIdentity]uint64)
	for _, sample := range records {
		if sample.presence != 0b11 {
			continue
		}
		total := sample.input + sample.output

		identity := tokenSourceIdentity{kind: sample.key.kind, name: sample.key.name}
		bucket := time.UnixMilli(sample.key.timestamp).Truncate(interval).Unix()
		byIdentity := bucketByIdentity[bucket]
		if byIdentity == nil {
			byIdentity = make(map[tokenSourceIdentity]uint64)
			bucketByIdentity[bucket] = byIdentity
		}
		byIdentity[identity] += total
		identityTotal[identity] += total
	}

	result.Trend = buildTokenSourceTrend(window, bucketByIdentity, identityTotal, now)
	return result
}

// buildTokenSourceTrend folds the per-bucket, per-identity totals into the
// bounded stacked-chart contract: Top N named identities (opaque keys) plus one
// global remainder, with one data point per aggregation bucket.
func buildTokenSourceTrend(
	window TimeWindow,
	bucketByIdentity map[int64]map[tokenSourceIdentity]uint64,
	identityTotal map[tokenSourceIdentity]uint64,
	now time.Time,
) TokenSourceTrend {
	interval := MachineDashboardAggregationInterval(window)
	identities := make([]tokenSourceIdentity, 0, len(identityTotal))
	for identity := range identityTotal {
		identities = append(identities, identity)
	}
	sort.Slice(identities, func(i, j int) bool {
		ti, tj := identityTotal[identities[i]], identityTotal[identities[j]]
		if ti == tj {
			if identities[i].kind == identities[j].kind {
				return identities[i].name < identities[j].name
			}
			return identities[i].kind < identities[j].kind
		}
		return ti > tj
	})
	top := identities
	hasRemainder := len(identities) > tokenSourceTrendTopN
	if hasRemainder {
		top = identities[:tokenSourceTrendTopN]
	}
	topKey := make(map[tokenSourceIdentity]string, len(top))
	for index, identity := range top {
		topKey[identity] = fmt.Sprintf("%s%d", tokenSourceTrendKeyPrefix, index)
	}
	series := make([]TokenSourceTrendSeries, 0, len(top)+1)
	for _, identity := range top {
		series = append(series, TokenSourceTrendSeries{
			Key: topKey[identity], Label: identity.name, SourceKind: identity.kind,
		})
	}
	if hasRemainder {
		series = append(series, TokenSourceTrendSeries{
			Key: tokenSourceTrendRemainderKey, Label: "Other",
		})
	}
	data := []TokenSourceTrendPoint{}
	if len(bucketByIdentity) > 0 {
		first := now.Add(-window.Duration()).Truncate(interval)
		last := now.Truncate(interval)
		for cursor := first; !cursor.After(last); cursor = cursor.Add(interval) {
			data = append(data, tokenSourceTrendPoint(cursor.Unix(), bucketByIdentity[cursor.Unix()], topKey, hasRemainder))
		}
	}
	return TokenSourceTrend{Series: series, Data: data}
}

func tokenSourceTrendPoint(
	timestamp int64,
	byIdentity map[tokenSourceIdentity]uint64,
	topKey map[tokenSourceIdentity]string,
	hasRemainder bool,
) TokenSourceTrendPoint {
	values := make(map[string]uint64, len(topKey)+1)
	var remainder uint64
	for identity, total := range byIdentity {
		if key, ok := topKey[identity]; ok {
			values[key] = total
		} else {
			remainder += total
		}
	}
	if hasRemainder {
		values[tokenSourceTrendRemainderKey] = remainder
	}
	return TokenSourceTrendPoint{Timestamp: timestamp, Values: values}
}

func normalizedTokenTrackingAgentIDs(agentIDs []string) []string {
	seen := make(map[string]struct{}, len(agentIDs))
	out := make([]string, 0, len(agentIDs))
	for _, value := range agentIDs {
		value = strings.TrimSpace(value)
		if value == "" || len(value) > 120 {
			continue
		}
		if _, exists := seen[value]; exists {
			continue
		}
		seen[value] = struct{}{}
		out = append(out, value)
	}
	sort.Strings(out)
	return out
}

func tokenTrackingValuesMatcher(labelName string, values []string) *labels.Matcher {
	if len(values) == 1 {
		return labels.MustNewMatcher(labels.MatchEqual, labelName, values[0])
	}
	escaped := make([]string, len(values))
	for i, value := range values {
		escaped[i] = regexp.QuoteMeta(value)
	}
	return labels.MustNewMatcher(labels.MatchRegexp, labelName, strings.Join(escaped, "|"))
}