package metrics

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

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

// QueryTimeSeries retrieves time-series data for the specified time window.
func (tc *TSDBCollector) QueryTimeSeries(timeWindow TimeWindow) []TimeSeriesDataPoint {
	tc.mu.RLock()
	defer tc.mu.RUnlock()

	if tc.db == nil {
		return []TimeSeriesDataPoint{}
	}

	interval := aggregationInterval(timeWindow)

	now := time.Now()
	windowDuration := timeWindow.Duration()
	mint := now.Add(-windowDuration).UnixMilli()
	maxt := now.UnixMilli()

	ctx := context.Background()
	querier, err := tc.db.Querier(mint, maxt)
	if err != nil {
		slog.Warn("Failed to create querier for time series", "error", err)
		return []TimeSeriesDataPoint{}
	}
	defer querier.Close()

	bucketMap := make(map[int64]*TimeSeriesDataPoint)

	// Query and aggregate each metric
	tc.aggregateTimeSeriesMetric(ctx, querier, metricRequestsTotal, interval, bucketMap, func(dp *TimeSeriesDataPoint, v float64) {
		dp.Count += uint64(v)
	})
	tc.aggregateTimeSeriesMetric(ctx, querier, metricInputTokensTotal, interval, bucketMap, func(dp *TimeSeriesDataPoint, v float64) {
		dp.Input += uint64(v)
	})
	tc.aggregateTimeSeriesMetric(ctx, querier, metricCachedInputTokensTotal, interval, bucketMap, func(dp *TimeSeriesDataPoint, v float64) {
		dp.Cached += uint64(v)
	})
	tc.aggregateTimeSeriesMetric(ctx, querier, metricOutputTokensTotal, interval, bucketMap, func(dp *TimeSeriesDataPoint, v float64) {
		dp.Output += uint64(v)
	})

	// Generate all buckets in the time window
	cutoff := now.Add(-windowDuration)
	firstBucket := cutoff.Truncate(interval)
	lastBucket := now.Truncate(interval)

	var result []TimeSeriesDataPoint
	for t := firstBucket; !t.After(lastBucket); t = t.Add(interval) {
		timestamp := t.Unix()
		dataPoint, exists := bucketMap[timestamp]
		if !exists {
			dataPoint = &TimeSeriesDataPoint{
				Timestamp: timestamp,
				Count:     0,
				Input:     0,
				Cached:    0,
				Output:    0,
			}
		}
		result = append(result, *dataPoint)
	}

	return result
}

// QueryTimeSeriesByKey retrieves time-series data for a specific API key.
func (tc *TSDBCollector) QueryTimeSeriesByKey(timeWindow TimeWindow, keyID string) []TimeSeriesDataPoint {
	tc.mu.RLock()
	defer tc.mu.RUnlock()

	if tc.db == nil {
		return []TimeSeriesDataPoint{}
	}

	interval := aggregationInterval(timeWindow)

	now := time.Now()
	windowDuration := timeWindow.Duration()
	mint := now.Add(-windowDuration).UnixMilli()
	maxt := now.UnixMilli()

	ctx := context.Background()
	querier, err := tc.db.Querier(mint, maxt)
	if err != nil {
		slog.Warn("Failed to create querier for time series by key", "error", err)
		return []TimeSeriesDataPoint{}
	}
	defer querier.Close()

	bucketMap := make(map[int64]*TimeSeriesDataPoint)

	// Query and aggregate each metric with keyID filter
	tc.aggregateTimeSeriesMetricWithFilter(ctx, querier, metricRequestsTotal, interval, bucketMap, func(dp *TimeSeriesDataPoint, v float64) {
		dp.Count += uint64(v)
	}, keyID, "")
	tc.aggregateTimeSeriesMetricWithFilter(ctx, querier, metricInputTokensTotal, interval, bucketMap, func(dp *TimeSeriesDataPoint, v float64) {
		dp.Input += uint64(v)
	}, keyID, "")
	tc.aggregateTimeSeriesMetricWithFilter(ctx, querier, metricCachedInputTokensTotal, interval, bucketMap, func(dp *TimeSeriesDataPoint, v float64) {
		dp.Cached += uint64(v)
	}, keyID, "")
	tc.aggregateTimeSeriesMetricWithFilter(ctx, querier, metricOutputTokensTotal, interval, bucketMap, func(dp *TimeSeriesDataPoint, v float64) {
		dp.Output += uint64(v)
	}, keyID, "")

	// Generate all buckets in the time window
	cutoff := now.Add(-windowDuration)
	firstBucket := cutoff.Truncate(interval)
	lastBucket := now.Truncate(interval)

	var result []TimeSeriesDataPoint
	for t := firstBucket; !t.After(lastBucket); t = t.Add(interval) {
		timestamp := t.Unix()
		dataPoint, exists := bucketMap[timestamp]
		if !exists {
			dataPoint = &TimeSeriesDataPoint{
				Timestamp: timestamp,
				Count:     0,
				Input:     0,
				Cached:    0,
				Output:    0,
			}
		}
		result = append(result, *dataPoint)
	}

	return result
}

// QueryTimeSeriesByModel retrieves time-series data for a specific model across all keys.
func (tc *TSDBCollector) QueryTimeSeriesByModel(timeWindow TimeWindow, model string) []TimeSeriesDataPoint {
	tc.mu.RLock()
	defer tc.mu.RUnlock()

	if tc.db == nil {
		return []TimeSeriesDataPoint{}
	}

	interval := aggregationInterval(timeWindow)

	now := time.Now()
	windowDuration := timeWindow.Duration()
	mint := now.Add(-windowDuration).UnixMilli()
	maxt := now.UnixMilli()

	ctx := context.Background()
	querier, err := tc.db.Querier(mint, maxt)
	if err != nil {
		slog.Warn("Failed to create querier for time series by model", "error", err)
		return []TimeSeriesDataPoint{}
	}
	defer querier.Close()

	bucketMap := make(map[int64]*TimeSeriesDataPoint)

	// Query and aggregate each metric with model filter
	tc.aggregateTimeSeriesMetricWithFilter(ctx, querier, metricRequestsTotal, interval, bucketMap, func(dp *TimeSeriesDataPoint, v float64) {
		dp.Count += uint64(v)
	}, "", model)
	tc.aggregateTimeSeriesMetricWithFilter(ctx, querier, metricInputTokensTotal, interval, bucketMap, func(dp *TimeSeriesDataPoint, v float64) {
		dp.Input += uint64(v)
	}, "", model)
	tc.aggregateTimeSeriesMetricWithFilter(ctx, querier, metricCachedInputTokensTotal, interval, bucketMap, func(dp *TimeSeriesDataPoint, v float64) {
		dp.Cached += uint64(v)
	}, "", model)
	tc.aggregateTimeSeriesMetricWithFilter(ctx, querier, metricOutputTokensTotal, interval, bucketMap, func(dp *TimeSeriesDataPoint, v float64) {
		dp.Output += uint64(v)
	}, "", model)

	// Generate all buckets in the time window
	cutoff := now.Add(-windowDuration)
	firstBucket := cutoff.Truncate(interval)
	lastBucket := now.Truncate(interval)

	var result []TimeSeriesDataPoint
	for t := firstBucket; !t.After(lastBucket); t = t.Add(interval) {
		timestamp := t.Unix()
		dataPoint, exists := bucketMap[timestamp]
		if !exists {
			dataPoint = &TimeSeriesDataPoint{
				Timestamp: timestamp,
				Count:     0,
				Input:     0,
				Cached:    0,
				Output:    0,
			}
		}
		result = append(result, *dataPoint)
	}

	return result
}

// QueryTimeSeriesByProvider retrieves time-series data for a specific provider across all keys.
// It matches all models with the "provider/" prefix in the model label.
func (tc *TSDBCollector) QueryTimeSeriesByProvider(timeWindow TimeWindow, provider string) []TimeSeriesDataPoint {
	tc.mu.RLock()
	defer tc.mu.RUnlock()

	if tc.db == nil {
		return []TimeSeriesDataPoint{}
	}

	interval := aggregationInterval(timeWindow)

	now := time.Now()
	windowDuration := timeWindow.Duration()
	mint := now.Add(-windowDuration).UnixMilli()
	maxt := now.UnixMilli()

	ctx := context.Background()
	querier, err := tc.db.Querier(mint, maxt)
	if err != nil {
		slog.Warn("Failed to create querier for time series by provider", "error", err)
		return []TimeSeriesDataPoint{}
	}
	defer querier.Close()

	bucketMap := make(map[int64]*TimeSeriesDataPoint)

	// Use regex to match all models with "provider/" prefix
	providerRegex := "^" + provider + "/.*"

	// Query and aggregate each metric with provider filter
	tc.aggregateTimeSeriesMetricByProvider(ctx, querier, metricRequestsTotal, interval, bucketMap, func(dp *TimeSeriesDataPoint, v float64) {
		dp.Count += uint64(v)
	}, providerRegex)
	tc.aggregateTimeSeriesMetricByProvider(ctx, querier, metricInputTokensTotal, interval, bucketMap, func(dp *TimeSeriesDataPoint, v float64) {
		dp.Input += uint64(v)
	}, providerRegex)
	tc.aggregateTimeSeriesMetricByProvider(ctx, querier, metricCachedInputTokensTotal, interval, bucketMap, func(dp *TimeSeriesDataPoint, v float64) {
		dp.Cached += uint64(v)
	}, providerRegex)
	tc.aggregateTimeSeriesMetricByProvider(ctx, querier, metricOutputTokensTotal, interval, bucketMap, func(dp *TimeSeriesDataPoint, v float64) {
		dp.Output += uint64(v)
	}, providerRegex)

	// Generate all buckets in the time window
	cutoff := now.Add(-windowDuration)
	firstBucket := cutoff.Truncate(interval)
	lastBucket := now.Truncate(interval)

	var result []TimeSeriesDataPoint
	for t := firstBucket; !t.After(lastBucket); t = t.Add(interval) {
		timestamp := t.Unix()
		dataPoint, exists := bucketMap[timestamp]
		if !exists {
			dataPoint = &TimeSeriesDataPoint{
				Timestamp: timestamp,
				Count:     0,
				Input:     0,
				Cached:    0,
				Output:    0,
			}
		}
		result = append(result, *dataPoint)
	}

	return result
}

// aggregateTimeSeriesMetricByProvider queries a metric filtered by provider regex on the model label,
// and aggregates values into time buckets.
func (tc *TSDBCollector) aggregateTimeSeriesMetricByProvider(
	ctx context.Context, querier storage.Querier, metricName string,
	aggregationInterval time.Duration, bucketMap map[int64]*TimeSeriesDataPoint,
	accumulator timeSeriesAccumulator, providerRegex string) {

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

	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(aggregationInterval)
			timestamp := bucketTime.Unix()

			existing, exists := bucketMap[timestamp]
			if !exists {
				existing = &TimeSeriesDataPoint{
					Timestamp: timestamp,
				}
				bucketMap[timestamp] = existing
			}
			accumulator(existing, v)
		}

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

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

type timeSeriesAccumulator func(*TimeSeriesDataPoint, float64)

// aggregateTimeSeriesMetric queries a metric and aggregates values into time buckets
func (tc *TSDBCollector) aggregateTimeSeriesMetric(
	ctx context.Context, querier storage.Querier, metricName string,
	aggregationInterval time.Duration, bucketMap map[int64]*TimeSeriesDataPoint, accumulator timeSeriesAccumulator) {

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

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

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

			existing, exists := bucketMap[timestamp]
			if !exists {
				existing = &TimeSeriesDataPoint{
					Timestamp: timestamp,
				}
				bucketMap[timestamp] = existing
			}
			accumulator(existing, v)
		}

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

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

// aggregateTimeSeriesMetricWithFilter queries a metric with optional keyID and model filters,
// and aggregates values into time buckets.
// For model filtering, it uses regex matching to support modelName (without provider prefix)
// matching against the full model identifier (provider/model format).
func (tc *TSDBCollector) aggregateTimeSeriesMetricWithFilter(
	ctx context.Context, querier storage.Querier, metricName string,
	aggregationInterval time.Duration, bucketMap map[int64]*TimeSeriesDataPoint,
	accumulator timeSeriesAccumulator, keyID, model string) {

	// Build matchers based on filters
	var matchers []*labels.Matcher
	matchers = append(matchers, labels.MustNewMatcher(labels.MatchEqual, "__name__", metricName))
	if keyID != "" {
		matchers = append(matchers, labels.MustNewMatcher(labels.MatchEqual, labelKeyID, keyID))
	}
	if model != "" {
		// Use regex matching to support both:
		// - Exact match: "gpt-4o" matches "gpt-4o"
		// - Provider prefix match: "gpt-4o" matches "openai/gpt-4o"
		// Regex matches either exact model name or "/model" at the end
		modelRegex := "^(|^.*/)" + model + "$"
		matchers = append(matchers, labels.MustNewMatcher(labels.MatchRegexp, labelModel, modelRegex))
	}

	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(aggregationInterval)
			timestamp := bucketTime.Unix()

			existing, exists := bucketMap[timestamp]
			if !exists {
				existing = &TimeSeriesDataPoint{
					Timestamp: timestamp,
				}
				bucketMap[timestamp] = existing
			}
			accumulator(existing, v)
		}

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

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