package central_thanos

import (
	"context"
	"fmt"
	"time"

	"github.com/tidbcloud/metrics-integration/pipeline/internal/datasource/client"
)

const nextGenStorageUsageQuantileWindow = "5m"

func buildNextGenRowStorageUsagePromQL(matcher string) string {
	return fmt.Sprintf(
		`quantile_over_time(0.5, sum(tikv_store_size_bytes{%s, type="used"})[%s]) or vector(0)`,
		matcher,
		nextGenStorageUsageQuantileWindow,
	)
}

func buildNextGenColumnarStorageUsagePromQL(matcher string) string {
	return fmt.Sprintf(
		`quantile_over_time(0.5, sum(tikv_store_size_bytes{%s, type="tiflash_used"})[%s]) or vector(0)`,
		matcher,
		nextGenStorageUsageQuantileWindow,
	)
}

// QueryDBRowStorageUsedBytes queries the row-based (TiKV) storage usage for a TiDB X cluster.
// Uses instant query when startTime == endTime, range query otherwise.
// clusterID corresponds to the per-plan cluster label in shared Thanos.
// The query preserves the current frontend semantics while avoiding a redundant
// shared-series regrouping step:
// 1. sum aggregates the row-storage gauge across all shared-series members
// 2. quantile_over_time(0.5, ...[5m]) smooths the storage signal over 5 minutes
// 3. or vector(0) ensures newly created clusters return 0 instead of empty
func QueryDBRowStorageUsedBytes(ctx context.Context, clusterID uint64, metricName string, startTime, endTime time.Time) (*client.Response, error) {
	matcher, err := sharedSeriesMatcher(ctx, clusterID, metricName, startTime)
	if err != nil {
		return nil, err
	}
	promql := buildNextGenRowStorageUsagePromQL(matcher)
	if startTime.Equal(endTime) {
		return client.DoInstantQueryWithTime(ctx, promql, metricName, startTime)
	}
	return client.DoRangeQuery(ctx, promql, metricName, startTime, endTime)
}

// QueryDBColumnarStorageUsedBytes queries the columnar (TiFlash) storage usage for a TiDB X cluster.
// Uses instant query when startTime == endTime, range query otherwise.
// Returns 0 when TiFlash is not enabled (via "or vector(0)" fallback).
// The query shape intentionally matches the current frontend storage chart semantics.
func QueryDBColumnarStorageUsedBytes(ctx context.Context, clusterID uint64, metricName string, startTime, endTime time.Time) (*client.Response, error) {
	matcher, err := sharedSeriesMatcher(ctx, clusterID, metricName, startTime)
	if err != nil {
		return nil, err
	}
	promql := buildNextGenColumnarStorageUsagePromQL(matcher)
	if startTime.Equal(endTime) {
		return client.DoInstantQueryWithTime(ctx, promql, metricName, startTime)
	}
	return client.DoRangeQuery(ctx, promql, metricName, startTime, endTime)
}