package metrics

import (
	"context"
	"log/slog"
	"sort"
	"time"

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

// RecordRequest records an access with model, token information, and duration.
func (tc *TSDBCollector) RecordRequest(
	apiKey, model string,
	inputToken, cachedInputToken, outputToken uint64,
	duration time.Duration) {

	keyID := KeyID(apiKey)
	timestamp := time.Now().UnixMilli()
	durationMs := float64(duration.Milliseconds())

	// Store masked key for display
	tc.mu.Lock()
	newKey := false
	if _, exists := tc.keyMaskMap[keyID]; !exists {
		tc.keyMaskMap[keyID] = MaskAPIKey(apiKey)
		newKey = true
	}
	tc.mu.Unlock()

	// Persist to disk if new key was added
	if newKey {
		if err := tc.saveKeyMaskMap(); err != nil {
			slog.Warn("Failed to save key mask map", "error", err)
		}
	}

	if tc.db == nil {
		slog.Warn("TSDB not available, skipping stats recording")
		return
	}

	ctx := context.Background()
	app := tc.db.Appender(ctx)

	baseLabels := labels.FromStrings(labelKeyID, keyID, labelModel, model)

	// Append request count
	if _, err := app.Append(0, buildLabels(baseLabels, metricRequestsTotal), timestamp, 1.0); err != nil {
		slog.Warn("Failed to append request count", "error", err)
	}

	// Append input tokens
	if _, err := app.Append(0, buildLabels(baseLabels, metricInputTokensTotal), timestamp, float64(inputToken)); err != nil {
		slog.Warn("Failed to append input tokens", "error", err)
	}

	// Append cached input tokens
	if _, err := app.Append(0, buildLabels(baseLabels, metricCachedInputTokensTotal), timestamp, float64(cachedInputToken)); err != nil {
		slog.Warn("Failed to append cached input tokens", "error", err)
	}

	// Append output tokens
	if outputToken > 0 {
		if _, err := app.Append(0, buildLabels(baseLabels, metricOutputTokensTotal), timestamp, float64(outputToken)); err != nil {
			slog.Warn("Failed to append output tokens", "error", err)
		}
	}

	// Append duration
	if durationMs > 0 {
		if _, err := app.Append(0, buildLabels(baseLabels, metricRequestDurationMs), timestamp, durationMs); err != nil {
			slog.Warn("Failed to append duration", "error", err)
		}
	}

	if err := app.Commit(); err != nil {
		slog.Warn("Failed to commit stats to TSDB", "error", err)
	}
}

// buildLabels creates Labels with the metric name added
func buildLabels(base labels.Labels, metricName string) labels.Labels {
	return labels.FromStrings(
		"__name__", metricName,
		labelKeyID, base.Get(labelKeyID),
		labelModel, base.Get(labelModel),
	)
}

func buildAgentTokenLabels(base labels.Labels, metricName string) labels.Labels {
	return labels.FromStrings(
		"__name__", metricName,
		labelAgentID, base.Get(labelAgentID),
		labelModel, base.Get(labelModel),
	)
}

// StatsSummary returns statistics summary for all API keys by querying TSDB.
func (tc *TSDBCollector) StatsSummary() map[string]*APIKeyStats {
	tc.mu.RLock()
	defer tc.mu.RUnlock()

	result := make(map[string]*APIKeyStats)

	if tc.db == nil {
		return result
	}

	ctx := context.Background()
	now := time.Now()
	mint := now.Add(-TimeWindow1Week.Duration()).UnixMilli()
	maxt := now.UnixMilli()

	querier, err := tc.db.Querier(mint, maxt)
	if err != nil {
		slog.Warn("Failed to create querier", "error", err)
		return result
	}
	defer querier.Close()

	// Collect all unique key_id and model combinations
	keyModelMap := tc.collectKeyModelCombinations(ctx, querier)

	// For each key/model combination, query and aggregate metrics
	for keyID, models := range keyModelMap {
		// Get masked key from map, fallback to short keyID if not found
		maskedKey := shortKeyID(keyID)
		tc.mu.RLock()
		if mk, exists := tc.keyMaskMap[keyID]; exists {
			maskedKey = mk
		}
		tc.mu.RUnlock()

		aks := &APIKeyStats{
			KeyID:      keyID,
			MaskedKey:  maskedKey,
			ModelStats: make(map[string]*ModelStats),
		}
		result[keyID] = aks

		for model := range models {
			modelStats := &ModelStats{Counters: make([]TimeCounter, 0)}
			aks.ModelStats[model] = modelStats

			tc.aggregateMetricsToCounters(ctx, querier, modelStats, keyID, model, now)
		}
	}

	return result
}

// collectKeyModelCombinations queries TSDB to find all unique key_id/model combinations
func (tc *TSDBCollector) collectKeyModelCombinations(ctx context.Context, querier storage.Querier) map[string]map[string]bool {
	keyModelMap := make(map[string]map[string]bool)

	matcher := labels.MustNewMatcher(labels.MatchEqual, "__name__", metricRequestsTotal)
	seriesSet := querier.Select(ctx, false, nil, matcher)

	for seriesSet.Next() {
		series := seriesSet.At()
		lbls := series.Labels()

		keyID := lbls.Get(labelKeyID)
		model := lbls.Get(labelModel)

		if keyID == "" || model == "" {
			continue
		}

		if keyModelMap[keyID] == nil {
			keyModelMap[keyID] = make(map[string]bool)
		}
		keyModelMap[keyID][model] = true
	}

	if err := seriesSet.Err(); err != nil {
		slog.Warn("Error iterating series", "error", err)
	}

	return keyModelMap
}

// aggregateMetricsToCounters queries all metrics for a key/model and aggregates into counters
func (tc *TSDBCollector) aggregateMetricsToCounters(
	ctx context.Context, querier storage.Querier, modelStats *ModelStats,
	keyID, model string, now time.Time) {

	bucketMap := make(map[time.Time]*TimeCounter)

	// Query and aggregate each metric
	tc.aggregateMetricByBucket(ctx, querier, metricRequestsTotal, keyID, model, bucketMap, func(c *TimeCounter, v float64) {
		c.Count += uint64(v)
	})
	tc.aggregateMetricByBucket(ctx, querier, metricInputTokensTotal, keyID, model, bucketMap, func(c *TimeCounter, v float64) {
		c.InputToken += uint64(v)
	})
	tc.aggregateMetricByBucket(ctx, querier, metricCachedInputTokensTotal, keyID, model, bucketMap, func(c *TimeCounter, v float64) {
		c.CachedInputToken += uint64(v)
	})
	tc.aggregateMetricByBucket(ctx, querier, metricOutputTokensTotal, keyID, model, bucketMap, func(c *TimeCounter, v float64) {
		c.OutputToken += uint64(v)
	})
	tc.aggregateMetricByBucket(ctx, querier, metricRequestDurationMs, keyID, model, bucketMap, func(c *TimeCounter, v float64) {
		c.Duration += uint64(v)
	})

	// Convert bucket map to sorted counters
	modelStats.mu.Lock()
	defer modelStats.mu.Unlock()

	for bucketTime, counter := range bucketMap {
		if counter.Count > 0 {
			modelStats.Counters = append(modelStats.Counters, TimeCounter{
				Timestamp:        bucketTime,
				Count:            counter.Count,
				InputToken:       counter.InputToken,
				CachedInputToken: counter.CachedInputToken,
				OutputToken:      counter.OutputToken,
				Duration:         counter.Duration,
			})
		}
	}

	// Sort counters by timestamp
	sort.Slice(modelStats.Counters, func(i, j int) bool {
		return modelStats.Counters[i].Timestamp.Before(modelStats.Counters[j].Timestamp)
	})
}

type counterAccumulator func(*TimeCounter, float64)

// aggregateMetricByBucket queries a single metric and aggregates values into 5-minute buckets
func (tc *TSDBCollector) aggregateMetricByBucket(
	ctx context.Context, querier storage.Querier, metricName, keyID, model string,
	bucketMap map[time.Time]*TimeCounter, accumulator counterAccumulator) {

	matchers := []*labels.Matcher{
		labels.MustNewMatcher(labels.MatchEqual, "__name__", metricName),
		labels.MustNewMatcher(labels.MatchEqual, labelKeyID, keyID),
		labels.MustNewMatcher(labels.MatchEqual, labelModel, model),
	}

	seriesSet := querier.Select(ctx, false, nil, matchers...)

	for seriesSet.Next() {
		series := seriesSet.At()
		iterator := series.Iterator(nil)

		for iterator.Next() == chunkenc.ValFloat {
			t, v := iterator.At()
			bucketTime := time.UnixMilli(t).Truncate(bucketGranularity)

			if bucketMap[bucketTime] == nil {
				bucketMap[bucketTime] = &TimeCounter{Timestamp: bucketTime}
			}
			accumulator(bucketMap[bucketTime], v)
		}

		if err := iterator.Err(); err != nil {
			slog.Warn("Error iterating samples", "error", err, "metric", metricName)
		}
	}

	if err := seriesSet.Err(); err != nil {
		slog.Warn("Error iterating series", "error", err, "metric", metricName)
	}
}