package shoot_prom

import (
	"context"
	"fmt"

	"github.com/pingcap/dbaas/central/metrics/prom"
	"github.com/tidbcloud/metrics/internal/datasource"
)

type MetricQuerierBuilder struct {
	client *prom.BasePromClient
}

func NewMetricQuerierBuilder(client *prom.BasePromClient) *MetricQuerierBuilder {
	return &MetricQuerierBuilder{
		client: client,
	}
}

func (b *MetricQuerierBuilder) BuildQueryRateQuerier() datasource.MetricQuerierFunc {
	return func(ctx context.Context) (prom.Metric, error) {
		// No need to further match cluster label, since we query the cluster internal Prometheus.
		vec, err := b.client.QueryVec(ctx, `sum by (type) (rate(tidb_executor_statement_total[1m]))`)
		if err != nil {
			return nil, err
		}
		r := prom.VecMetric{
			Name:    "tidbcloud_db_query_rate",
			Help:    "Queries per second by statement type",
			Type:    "gauge",
			Metrics: make([]*prom.VecMetricElement, 0, vec.Len()),
		}
		for _, m := range vec {
			r.Metrics = append(r.Metrics, &prom.VecMetricElement{
				Labels: []prom.MetricLabel{{Name: "sql_type", Value: string(m.Metric["type"])}},
				Value:  m.Value,
			})
		}
		return &r, nil
	}
}

func (b *MetricQuerierBuilder) BuildFailedQueryRateQuerier() datasource.MetricQuerierFunc {
	return func(ctx context.Context) (prom.Metric, error) {
		vec, err := b.client.QueryVec(ctx, `sum by (type) (rate(tidb_server_execute_error_total[1m]))`)
		if err != nil {
			return nil, err
		}
		r := prom.VecMetric{
			Name:    "tidbcloud_db_failed_query_rate",
			Help:    "Failed queries per second",
			Type:    "gauge",
			Metrics: make([]*prom.VecMetricElement, 0, vec.Len()),
		}
		for _, m := range vec {
			r.Metrics = append(r.Metrics, &prom.VecMetricElement{
				Labels: []prom.MetricLabel{{Name: "type", Value: string(m.Metric["type"])}},
				Value:  m.Value,
			})
		}
		return &r, nil
	}
}

func (b *MetricQuerierBuilder) BuildQueryDurationQuerier() datasource.MetricQuerierFunc {
	return func(ctx context.Context) (prom.Metric, error) {
		queryTemplate := `
		histogram_quantile(%s,
		  sum by (sql_type, le) (rate(tidb_server_handle_query_duration_seconds_bucket[1m]))
		)
	`
		quantiles := []string{"0.99", "0.80"}
		queries := make([]string, 0)
		for _, q := range quantiles {
			queries = append(queries, fmt.Sprintf(queryTemplate, q))
		}

		vecs, err := b.client.MultiQueryVec(ctx, queries...)
		if err != nil {
			return nil, err
		}
		r := prom.VecMetric{
			Name:    "tidbcloud_db_query_duration_seconds",
			Help:    "Queries durations of different quantiles aggregated in last minute by statement type",
			Type:    "gauge",
			Metrics: make([]*prom.VecMetricElement, 0, len(quantiles)*10),
		}
		for idx, vec := range vecs {
			for _, m := range vec {
				r.Metrics = append(r.Metrics, &prom.VecMetricElement{
					Labels: []prom.MetricLabel{
						{Name: "quantile", Value: quantiles[idx]},
						{Name: "sql_type", Value: string(m.Metric["sql_type"])},
					},
					Value: m.Value,
				})
			}
		}
		return &r, nil
	}
}

func (b MetricQuerierBuilder) BuildCPUUsageCoresQuerier() datasource.MetricQuerierFunc {
	return func(ctx context.Context) (prom.Metric, error) {
		vec, err := b.client.QueryVec(ctx, `sum by (instance, component)(rate(process_cpu_seconds_total{component=~"tidb|tikv|tiflash"}[1m]))`)
		if err != nil {
			return nil, err
		}
		r := prom.VecMetric{
			Name:    "tidbcloud_process_cpu_seconds_rate",
			Help:    "CPU usage percent by instances",
			Type:    "gauge",
			Metrics: make([]*prom.VecMetricElement, 0, vec.Len()),
		}
		for _, m := range vec {
			r.Metrics = append(r.Metrics, &prom.VecMetricElement{
				Labels: []prom.MetricLabel{
					{Name: datasource.InstanceName, Value: datasource.UnifyDBName(string(m.Metric["instance"]))},
					{Name: datasource.ComponentType, Value: datasource.UnifyDBName(string(m.Metric["component"]))},
				},
				Value: m.Value,
			})
		}
		return &r, nil
	}
}

func (b MetricQuerierBuilder) BuildConnectionsQuerier() datasource.MetricQuerierFunc {
	return func(ctx context.Context) (prom.Metric, error) {
		vec, err := b.client.QueryVec(ctx, `sum(tidb_server_connections)`)
		if err != nil {
			return nil, err
		}
		r := prom.SingleMetric{
			Name:  "tidbcloud_db_connection_count",
			Help:  "Current connection count",
			Type:  "gauge",
			Value: 0,
		}
		if len(vec) > 0 {
			r.Value = vec[0].Value
		}
		return &r, nil
	}
}

func (b MetricQuerierBuilder) BuildStorageUsageBytesQuerier() datasource.MetricQuerierFunc {
	return func(ctx context.Context) (prom.Metric, error) {
		vec, err := b.client.QueryVec(ctx, `sum(pd_cluster_status{type="storage_size"})`)
		if err != nil {
			return nil, err
		}
		r := prom.SingleMetric{
			Name:  "tidbcloud_storage_used_bytes",
			Help:  "Current used storage size in bytes",
			Type:  "gauge",
			Value: 0,
		}
		if len(vec) > 0 {
			r.Value = vec[0].Value
		}
		return &r, nil
	}
}

func (b MetricQuerierBuilder) BuildStorageCapacityBytesQuerier() datasource.MetricQuerierFunc {
	return func(ctx context.Context) (prom.Metric, error) {
		vec, err := b.client.QueryVec(ctx, `sum(pd_cluster_status{type="storage_capacity"})`)
		if err != nil {
			return nil, err
		}
		r := prom.SingleMetric{
			Name:  "tidbcloud_storage_capacity_bytes",
			Help:  "Current storage capacity in bytes",
			Type:  "gauge",
			Value: 0,
		}
		if len(vec) > 0 {
			r.Value = vec[0].Value
		}
		return &r, nil
	}
}