package thanos

import (
	"context"
	"fmt"
	"strconv"
	"strings"
	"sync"
	"time"

	"github.com/prometheus/client_golang/prometheus"
	"github.com/tidbcloud/metrics-server/central/db/entity"
	"github.com/tidbcloud/metrics-server/internal/datasource"
	"github.com/tidbcloud/metrics-server/internal/datasource/central_thanos"
	"github.com/tidbcloud/metrics-server/internal/datasource/client"
	"github.com/tidbcloud/metrics-server/internal/metadata"
	"github.com/tidbcloud/metrics-server/internal/otel_collector"
	"go.opentelemetry.io/collector/component"
	otelconfig "go.opentelemetry.io/collector/config"
	"go.opentelemetry.io/collector/consumer"
	"go.opentelemetry.io/collector/exporter/exporterhelper"
	"go.opentelemetry.io/collector/pdata/pcommon"
	"go.opentelemetry.io/collector/pdata/pmetric"
	"go.opentelemetry.io/collector/processor/processorhelper"
	"go.uber.org/zap"
)

const (
	typeStr = "thanosquery"
)

var (
	consumerCapabilities = consumer.Capabilities{MutatesData: true}
	// Same as the metric in tidb: https://github.com/pingcap/tidb/blob/master/metrics/server.go#L45.
	durationBuckets = prometheus.ExponentialBuckets(0.0005, 2, 29)

	processingDuration = prometheus.NewHistogramVec(
		prometheus.HistogramOpts{
			Name:    "otelcol_processor_query_duration",
			Help:    "the time that used for querying metrics from o11y metrics backend",
			Buckets: prometheus.ExponentialBuckets(0.5, 2, 10), // 500ms ~ 512s
		}, []string{"integration_type"})
)

func init() {
	prometheus.MustRegister(processingDuration)
}

type Config struct {
	otelconfig.ProcessorSettings `mapstructure:",squash"`
	// Currently no usage.
	exporterhelper.TimeoutSettings `mapstructure:",squash"`
}

func NewFactory() component.ProcessorFactory {
	return component.NewProcessorFactory(
		typeStr,
		createDefaultConfig,
		component.WithMetricsProcessor(createMetricsProcessor, component.StabilityLevelStable))
}

func createDefaultConfig() otelconfig.Processor {
	return &Config{
		ProcessorSettings: otelconfig.NewProcessorSettings(otelconfig.NewComponentID(typeStr)),
		TimeoutSettings:   exporterhelper.TimeoutSettings{Timeout: 15 * time.Second},
	}
}

func createMetricsProcessor(
	_ context.Context,
	params component.ProcessorCreateSettings,
	pConf otelconfig.Processor,
	nextConsumer consumer.Metrics,
) (component.MetricsProcessor, error) {
	cfg := pConf.(*Config)
	thanosQuerier, err := newCentralThanosQuerier(params.Logger, cfg)
	if err != nil {
		return nil, err
	}
	return processorhelper.NewMetricsProcessor(cfg, nextConsumer, thanosQuerier.doQuery, processorhelper.WithCapabilities(consumerCapabilities))
}

type centralThanosQuerier struct {
	config *Config

	logger *zap.Logger
}

func newCentralThanosQuerier(logger *zap.Logger, config *Config) (*centralThanosQuerier, error) {
	querier := centralThanosQuerier{
		config: config,
		logger: logger,
	}
	return &querier, nil
}

// This function assumes the upstream receiver can only be metadb receiver.
func (q *centralThanosQuerier) doQuery(ctx context.Context, md pmetric.Metrics) (pmetric.Metrics, error) {
	start := time.Now()
	wg := &sync.WaitGroup{}
	// integrationType will be updated during
	integrationType := ""
	// we consider integrationType for all resourceMetrics are the same in a single doQuery
	if md.ResourceMetrics().Len() > 0 {
		if _integrationType, ok := md.ResourceMetrics().At(0).Resource().Attributes().Get(otel_collector.IntegrationType); ok {
			integrationType = _integrationType.StringVal()
		}
	}
	defer func() {
		processingDuration.
			WithLabelValues(integrationType).Observe(time.Since(start).Seconds())
	}()

	for i := 0; i < md.ResourceMetrics().Len(); i++ {
		// Reuse the ResourceMetric object generated by the upstream metadb receiver.
		rm := md.ResourceMetrics().At(i)
		// In the ResourceMetric object, create a different InstrumentationLibraryMetric object.
		ilm := rm.ScopeMetrics().AppendEmpty()
		ilm.Scope().SetName("otel/tidbcloud/thanos")
		now := pcommon.NewTimestampFromTime(time.Now())
		mem := ilm.Metrics().AppendEmpty()
		metadata.M.TidbCloudNodeMemoryUsedBytes.Init(mem)
		cpu := ilm.Metrics().AppendEmpty()
		metadata.M.TidbCloudNodeCPUSecondsTotal.Init(cpu)
		disk := ilm.Metrics().AppendEmpty()
		metadata.M.TidbCloudNodeStorageUsedBytes.Init(disk)
		conns := ilm.Metrics().AppendEmpty()
		metadata.M.TidbCloudDbConnections.Init(conns)
		queries := ilm.Metrics().AppendEmpty()
		metadata.M.TidbCloudDbQueriesTotal.Init(queries)
		fails := ilm.Metrics().AppendEmpty()
		metadata.M.TidbCloudDbFailedQueriesTotal.Init(fails)
		duration := ilm.Metrics().AppendEmpty()
		metadata.M.TidbCloudDbQueryDurationSeconds.Init(duration)
		totalConn := ilm.Metrics().AppendEmpty()
		metadata.M.TidbCloudDbTotalConnection.Init(totalConn)
		activeConns := ilm.Metrics().AppendEmpty()
		metadata.M.TidbCloudDbActiveConnections.Init(activeConns)
		disconns := ilm.Metrics().AppendEmpty()
		metadata.M.TidbCloudDbDisconnections.Init(disconns)
		dbTime := ilm.Metrics().AppendEmpty()
		metadata.M.TidbCloudDbDatabaseTime.Init(dbTime)
		qps := ilm.Metrics().AppendEmpty()
		metadata.M.TidbCloudDbQueryPerSecond.Init(qps)
		failedQueries := ilm.Metrics().AppendEmpty()
		metadata.M.TidbCloudDbFailedQueries.Init(failedQueries)
		cmdPerSecond := ilm.Metrics().AppendEmpty()
		metadata.M.TidbCloudDbCommandPerSecond.Init(cmdPerSecond)
		queriesUsingPlanCache := ilm.Metrics().AppendEmpty()
		metadata.M.TidbCloudDbQueriesUsingPlanCacheOps.Init(queriesUsingPlanCache)
		avgQueryDuration := ilm.Metrics().AppendEmpty()
		metadata.M.TidbCloudDbAverageQueryDuration.Init(avgQueryDuration)
		tps := ilm.Metrics().AppendEmpty()
		metadata.M.TidbCloudDbTransactionPerSecond.Init(tps)

		// Access meta attributes generated by upstream metadb receiver.
		attrs := rm.Resource().Attributes()
		_id, ok := attrs.Get(otel_collector.ClusterIDKey)
		if !ok {
			// Log then go to the next cluster id.
			q.logger.Warn("failed to get cluster id")
			continue
		}
		id := _id.StringVal()
		_name, ok := attrs.Get(otel_collector.ClusterNameKey)
		if !ok {
			q.logger.Warn("failed to get cluster name")
			continue
		}
		name := _name.StringVal()
		clusterID, err := strconv.ParseUint(id, 10, 64)
		if err != nil {
			// Log then go to the next cluster id.
			q.logger.Warn("Failed to parse cluster id", zap.Error(err))
			continue
		}
		_projname, ok := attrs.Get(otel_collector.ProjectNameKey)
		if !ok {
			q.logger.Warn("failed to get project name")
			continue
		}
		projname := _projname.StringVal()

		ctx, err := datasource.SetMetricDataAddressToContextWithCache(ctx, entity.ID(clusterID))
		if err != nil {
			// Log then go to the next cluster id.
			q.logger.Warn("Failed to call datasource.SetMetricDataAddressToContextWithCache", zap.Error(err))
			continue
		}

		// Prepare the extra label apply functions.
		doNothing := func(pcommon.Map, client.LabelSet) {}
		applyErrTypeLabel := func(attrs pcommon.Map, ls client.LabelSet) {
			attrs.Insert(metadata.A.ErrorType, pcommon.NewValueString(ls.Type))
		}
		applySqlTypeLabel := func(attrs pcommon.Map, ls client.LabelSet) {
			attrs.Insert(metadata.A.SQLType, pcommon.NewValueString(ls.Type))
		}
		applySqlTypeLabelV2 := func(attrs pcommon.Map, ls client.LabelSet) {
			attrs.Insert(metadata.A.SQLType, pcommon.NewValueString(ls.SqlType))
		}
		applyResultLabel := func(attrs pcommon.Map, ls client.LabelSet) {
			attrs.Insert(metadata.A.Result, pcommon.NewValueString(ls.Result))
		}
		applyLabelsForTransaction := func(attrs pcommon.Map, ls client.LabelSet) {
			attrs.Insert(metadata.A.TxnMode, pcommon.NewValueString(ls.TxnMode))
			attrs.Insert(metadata.A.ErrorType, pcommon.NewValueString(ls.Type))
		}

		// Parallelize processing for each cluster.
		wg.Add(1)
		go func() {
			defer wg.Done()
			// Query TiDB cluster metrics one by one.
			if localErr := q.queryThenWriteNumMetric(mem, central_thanos.QueryNodeMemoryUsedBytes, setDoubleMetric, ctx, clusterID, name, projname, now, doNothing); localErr != nil {
				// Log failure then go to the next query.
				q.logger.Warn("QueryNodeMemoryUsedBytes failed", zap.Error(localErr))
			}
			if localErr := q.queryThenWriteNumMetric(cpu, central_thanos.QueryNodeCPUUsedSeconds, setDoubleMetric, ctx, clusterID, name, projname, now, doNothing); localErr != nil {
				q.logger.Warn("QueryNodeCPUUsedSeconds failed", zap.Error(localErr))
			}
			if localErr := q.queryThenWriteNumMetric(disk, central_thanos.QueryNodeStorageUsedBytes, setIntMetric, ctx, clusterID, name, projname, now, doNothing); localErr != nil {
				q.logger.Warn("QueryNodeStorageUsedBytes failed", zap.Error(localErr))
			}
			if localErr := q.queryThenWriteNumMetric(conns, central_thanos.QueryDBConnections, setIntMetric, ctx, clusterID, name, projname, now, doNothing); localErr != nil {
				q.logger.Warn("QueryDBConnections failed", zap.Error(localErr))
			}
			if localErr := q.queryThenWriteNumMetric(queries, central_thanos.QueryDBQueriesTotal, setDoubleMetric, ctx, clusterID, name, projname, now, applySqlTypeLabel); localErr != nil {
				q.logger.Warn("QueryDBQueriesTotal failed", zap.Error(localErr))
			}
			if localErr := q.queryThenWriteNumMetric(fails, central_thanos.QueryDBFailedQueriesTotal, setIntMetric, ctx, clusterID, name, projname, now, applyErrTypeLabel); localErr != nil {
				q.logger.Warn("QueryDBFailedQueriesTotal failed", zap.Error(localErr))
			}
			if localErr := q.queryThenWriteDurationHistogramMetric(duration, central_thanos.QueryDBQueryDurationSeconds, ctx, clusterID, name, projname, now, applySqlTypeLabelV2); localErr != nil {
				q.logger.Warn("QueryDBQueryDurationSeconds failed", zap.Error(localErr))
			}
			if localErr := q.queryThenWriteNumMetric(totalConn, central_thanos.QueryDBTotalConnection, setDoubleMetric, ctx, clusterID, name, projname, now, doNothing); localErr != nil {
				q.logger.Warn("QueryDBTotalConnection failed", zap.Error(localErr))
			}
			if localErr := q.queryThenWriteNumMetric(activeConns, central_thanos.QueryDBActiveConnections, setDoubleMetric, ctx, clusterID, name, projname, now, doNothing); localErr != nil {
				q.logger.Warn("QueryDBActiveConnections failed", zap.Error(localErr))
			}
			if localErr := q.queryThenWriteNumMetric(disconns, central_thanos.QueryDBDisconnections, setDoubleMetric, ctx, clusterID, name, projname, now, applyResultLabel); localErr != nil {
				q.logger.Warn("QueryDBDisconnections failed", zap.Error(localErr))
			}
			if localErr := q.queryThenWriteNumMetric(dbTime, central_thanos.QueryDBDatabaseTime, setDoubleMetric, ctx, clusterID, name, projname, now, applySqlTypeLabelV2); localErr != nil {
				q.logger.Warn("QueryDBDatabaseTime failed", zap.Error(localErr))
			}
			if localErr := q.queryThenWriteNumMetric(qps, central_thanos.QueryDBQueryPerSecond, setDoubleMetric, ctx, clusterID, name, projname, now, applyErrTypeLabel); localErr != nil {
				q.logger.Warn("QueryDBQueryPerSecond failed", zap.Error(localErr))
			}
			if localErr := q.queryThenWriteNumMetric(failedQueries, central_thanos.QueryDBFailedQueries, setDoubleMetric, ctx, clusterID, name, projname, now, applyErrTypeLabel); localErr != nil {
				q.logger.Warn("QueryDBFailedQueries failed", zap.Error(localErr))
			}
			if localErr := q.queryThenWriteNumMetric(cmdPerSecond, central_thanos.QueryDBCommandPerSecond, setDoubleMetric, ctx, clusterID, name, projname, now, applyErrTypeLabel); localErr != nil {
				q.logger.Warn("QueryDBCommandPerSecond failed", zap.Error(localErr))
			}
			if localErr := q.queryThenWriteNumMetric(queriesUsingPlanCache, central_thanos.QueryDBQueriesUsingPlanCacheOps, setDoubleMetric, ctx, clusterID, name, projname, now, doNothing); localErr != nil {
				q.logger.Warn("QueryDBQueriesUsingPlanCacheOps failed", zap.Error(localErr))
			}
			if localErr := q.queryThenWriteNumMetric(avgQueryDuration, central_thanos.QueryDBAverageQueryDuration, setDoubleMetric, ctx, clusterID, name, projname, now, applySqlTypeLabelV2); localErr != nil {
				q.logger.Warn("QueryDBAverageQueryDuration failed", zap.Error(localErr))
			}
			if localErr := q.queryThenWriteNumMetric(tps, central_thanos.QueryDBTransactionPerSecond, setDoubleMetric, ctx, clusterID, name, projname, now, applyLabelsForTransaction); localErr != nil {
				q.logger.Warn("QueryDBTransactionPerSecond failed", zap.Error(localErr))
			}
		}()
	}
	wg.Wait()

	return md, nil
}

type extraLabelFunc func(pcommon.Map, client.LabelSet)

func (q *centralThanosQuerier) queryThenWriteNumMetric(
	m pmetric.Metric,
	queryFunc func(ctx context.Context, clusterID uint64, metric string) (*client.Response, error),
	setNumFunc func(p pmetric.NumberDataPoint, v string) error,
	ctx context.Context,
	clusterID uint64,
	clusterName string,
	projectName string,
	t pcommon.Timestamp,
	extraLabelFunc extraLabelFunc) error {

	res, err := queryFunc(ctx, clusterID, m.Name())
	if err != nil || res.Status != "success" {
		return fmt.Errorf("failed to call queryFunc, err: %v, result: %+v", err, res)
	}

	for _, series := range res.Data.Result {
		// Parse results from the query request.
		labels, value := series.LabelSet, series.Value
		if len(value) != 2 {
			return fmt.Errorf("prometheus response failed: value(%+v)", value)
		}

		// No need to get pd metrics.
		if labels.Component == "pd" {
			continue
		}

		// Build labels.
		attrs := pcommon.NewMap()
		attrs.Insert(metadata.A.ProjectName, pcommon.NewValueString(projectName))
		attrs.Insert(metadata.A.ClusterName, pcommon.NewValueString(clusterName))
		attrs.Insert(metadata.A.ClusterID, pcommon.NewValueString(strconv.FormatUint(clusterID, 10)))
		attrs.Insert(metadata.A.Component, pcommon.NewValueString(labels.Component))
		attrs.Insert(metadata.A.Instance, pcommon.NewValueString(datasource.UnifyDBName(labels.Instance)))
		extraLabelFunc(attrs, labels)

		// Write the data point and labels.
		var p pmetric.NumberDataPoint
		switch m.DataType() {
		case pmetric.MetricDataTypeGauge:
			p = m.Gauge().DataPoints().AppendEmpty()
		case pmetric.MetricDataTypeSum:
			p = m.Sum().DataPoints().AppendEmpty()
		default:
			return fmt.Errorf("unsupported data type: %v", m.DataType())
		}
		p.SetTimestamp(t)
		if err := setNumFunc(p, value[1].(string)); err != nil {
			return err
		}
		attrs.CopyTo(p.Attributes())
	}

	return nil
}

func (q *centralThanosQuerier) queryThenWriteDurationHistogramMetric(
	m pmetric.Metric,
	queryFunc func(ctx context.Context, clusterID uint64, metric string) (*client.Response, error),
	ctx context.Context,
	clusterID uint64,
	clusterName string,
	projectName string,
	t pcommon.Timestamp,
	extraLabelFunc extraLabelFunc) error {

	res, err := queryFunc(ctx, clusterID, m.Name())
	if err != nil || res.Status != "success" {
		return fmt.Errorf("failed to call queryFunc, err: %v, result: %+v", err, res)
	}

	histoMap := q.calculateBuckets(res.Data.Result)

	for desc, data := range histoMap {
		// No need to get pd metrics.
		if desc.Component == "pd" {
			continue
		}

		attrs := pcommon.NewMap()
		attrs.Insert(metadata.A.ProjectName, pcommon.NewValueString(projectName))
		attrs.Insert(metadata.A.ClusterID, pcommon.NewValueString(strconv.FormatUint(clusterID, 10)))
		attrs.Insert(metadata.A.ClusterName, pcommon.NewValueString(clusterName))
		attrs.Insert(metadata.A.Component, pcommon.NewValueString(desc.Component))
		attrs.Insert(metadata.A.Instance, pcommon.NewValueString(datasource.UnifyDBName(desc.Instance)))
		extraLabelFunc(attrs, desc)

		p := m.Histogram().DataPoints().AppendEmpty()
		p.SetTimestamp(t)
		p.SetCount(data.count)
		p.SetSum(data.sum)
		p.SetBucketCounts(pcommon.NewImmutableUInt64Slice(data.bucketCounts))
		p.SetExplicitBounds(pcommon.NewImmutableFloat64Slice(durationBuckets))
		attrs.CopyTo(p.Attributes())
	}

	return nil
}

func setIntMetric(p pmetric.NumberDataPoint, v string) error {
	num, err := strconv.ParseInt(v, 10, 64)
	if err != nil {
		f64, err := strconv.ParseFloat(v, 64)
		if err != nil {
			return fmt.Errorf("failed to parse value: %v", v)
		}
		num = int64(f64)
	}
	p.SetIntVal(num)
	return nil
}

func setDoubleMetric(p pmetric.NumberDataPoint, v string) error {
	num, err := strconv.ParseFloat(v, 64)
	if err != nil {
		return fmt.Errorf("failed to parse value: %v", v)
	}
	p.SetDoubleVal(num)
	return nil
}

type histoData struct {
	sum          float64
	count        uint64
	bucketCounts []uint64
}

// calculateBuckets group raw series by histoDesc.
// It assumes results contains _sum, _count and _bucket series,
// And _bucket series contain the le label.
func (q *centralThanosQuerier) calculateBuckets(res []client.Result) map[client.LabelSet]*histoData {
	histoMap := make(map[client.LabelSet]*histoData)
	for _, series := range res {
		// For each series, build a LabelSet as the group key.
		desc := client.LabelSet{
			// KubernetesNamespace contains cluster id.
			KubernetesNamespace: series.KubernetesNamespace,
			Component:           series.Component,
			Instance:            series.Instance,
			SqlType:             series.SqlType,
			Type:                series.Type,
			InTxn:               series.InTxn,
			TxnMode:             series.TxnMode,
		}

		if _, ok := histoMap[desc]; !ok {
			histoMap[desc] = &histoData{
				sum:          0,
				count:        0,
				bucketCounts: make([]uint64, len(durationBuckets)+1),
			}
		}
		data := histoMap[desc]

		// Check series data point.
		if len(series.Value) != 2 {
			// Go to the next series.
			q.logger.Warn("calculateBuckets: prometheus response failed", zap.Any("value", series.Value))
			continue
		}

		// Reduce series to histoData.
		if strings.HasSuffix(series.MetricName, "_sum") {
			num, err := strconv.ParseFloat(series.Value[1].(string), 64)
			if err != nil {
				// Go to the next series.
				q.logger.Warn("calculateBuckets: parsing _sum failed", zap.Any("value", series.Value))
				continue
			}
			data.sum = num
		} else if strings.HasSuffix(series.MetricName, "_count") {
			num, err := strconv.ParseFloat(series.Value[1].(string), 64)
			if err != nil {
				q.logger.Warn("calculateBuckets: parsing _count failed", zap.Any("value", series.Value))
				continue
			}
			data.count = uint64(num)
		} else if strings.HasSuffix(series.MetricName, "_bucket") {
			le, err := strconv.ParseFloat(series.Le, 64)
			if err != nil {
				q.logger.Warn("calculateBuckets: parsing _bucket le failed", zap.String("le", series.Le))
				continue
			}
			num, err := strconv.ParseFloat(series.Value[1].(string), 64)
			if err != nil {
				q.logger.Warn("calculateBuckets: parsing _bucket value failed", zap.Any("value", series.Value))
				continue
			}
			data.bucketCounts[findBucketIdx(le)] = uint64(num)
		}
	}

	// Caution: Buckets in prometheus are cumulative but buckets in otel are not!
	for _, data := range histoMap {
		// The first bucket is not cumulative.
		adjusted := make([]uint64, len(data.bucketCounts))
		adjusted[0] = data.bucketCounts[0]
		for i := 1; i < len(data.bucketCounts); i++ {
			adjusted[i] = data.bucketCounts[i] - data.bucketCounts[i-1]
		}
		data.bucketCounts = adjusted
	}

	return histoMap
}

// findBucketIdx find the position in bucketCounts for a given upper bound.
func findBucketIdx(le float64) int {
	for i, f := range durationBuckets {
		if le > f {
			continue
		}
		return i
	}
	return len(durationBuckets)
}