package metadb

import (
	"context"

	"github.com/pingcap/dbaas/central/db"
	"github.com/pingcap/dbaas/central/db/dbutil"
	"github.com/pingcap/dbaas/central/db/entity"
	"github.com/pingcap/dbaas/central/metrics/prom"
	clustermodel "github.com/pingcap/dbaas/central/model/cluster"
	"github.com/pingcap/dbaas/pkg/tidb/component"
	"github.com/pingcap/errors"
	"github.com/prometheus/common/model"
	"github.com/tidbcloud/metrics/internal/datasource"
	"github.com/tidbcloud/metrics/internal/datasource/central_thanos"
	"gorm.io/gorm"
)

func BuildNodeMemoryCapacityQuerier(clusterID uint64, useClusterMetrics bool) datasource.MetricQuerierFunc {
	return func(ctx context.Context) (prom.Metric, error) {
		vec := prom.VecMetric{
			Name:    "tidbcloud_node_memory_capacity_bytes",
			Help:    "The memory capacity bytes of TiDB/TiKV/TiFlash nodes",
			Type:    "gauge",
			Metrics: []*prom.VecMetricElement{},
		}

		var (
			err                     error
			componentResources      []component.Resources
			componentCapacityMapper = map[string]uint64{}
			cluster                 entity.Cluster
		)

		// Get memory capacity from the `machine_profiles` table in the metadb.
		err = db.Transaction(func(tx *gorm.DB) error {
			cluster, err = dbutil.GetCluster(tx, entity.ID(clusterID))
			if err != nil {
				return err
			}
			componentResources, err = clustermodel.GetClusterComponentsV2(cluster)
			if err != nil {
				return err
			}
			return nil
		})
		if err != nil {
			return &vec, err
		}

		for _, componentResource := range componentResources {
			if componentResource.ComponentName.String() == component.PD.String() {
				continue
			}
			componentName := componentResource.ComponentName.String()
			componentCapacityMapper[componentName] = uint64(componentResource.Capacity.Compute.MemoryMi) * 1024 * 1024
		}

		// Check id against the hardcoded whitelist.
		if err := datasource.IsOKToScrapeNodeMemory(clusterID); err != nil {
			return &vec, err
		}

		// We can only get instance from thanos metrics.
		thanosNodeMemoryRespResult, err := central_thanos.QueryNodeMemoryUsedBytes(ctx, clusterID)
		if err != nil || thanosNodeMemoryRespResult.Status != "success" {
			return &vec, errors.Errorf("failed to query used node memory, status: %s", thanosNodeMemoryRespResult.Status)
		}

		for _, result := range thanosNodeMemoryRespResult.Data.Result {
			m := result.LabelSet
			valueTuple := result.Value
			if len(valueTuple) != 2 {
				return &vec, errors.Errorf("failed to query value(%+v) of used node memory", valueTuple)
			}

			if m.Component == component.PD.String() {
				continue
			}

			capacityValue := componentCapacityMapper[m.Component]
			var labels []prom.MetricLabel
			if useClusterMetrics {
				labels = []prom.MetricLabel{
					{Name: datasource.LabelInstance, Value: datasource.UnifyDBName(m.Instance)},
					{Name: datasource.LabelComponent, Value: datasource.UnifyDBName(m.Component)},
					{Name: datasource.LabelCluster, Value: cluster.Name},
				}
			} else {
				labels = []prom.MetricLabel{
					{Name: datasource.InstanceName, Value: datasource.UnifyDBName(m.Instance)},
					{Name: datasource.ComponentType, Value: datasource.UnifyDBName(m.Component)},
				}
			}

			vec.Metrics = append(vec.Metrics, &prom.VecMetricElement{
				Labels: labels,
				Value:  model.SampleValue(capacityValue),
			})
		}

		return &vec, nil
	}

}

func BuildNodeCPUCapacityCoresQuerier(clusterID uint64) datasource.MetricQuerierFunc {
	return func(ctx context.Context) (prom.Metric, error) {
		vec := prom.VecMetric{
			Name:    "tidbcloud_node_cpu_capacity_cores",
			Help:    "The CPU limit cores of TiDB/TiKV/TiFlash nodes",
			Type:    "gauge",
			Metrics: []*prom.VecMetricElement{},
		}

		var (
			err                     error
			componentResources      []component.Resources
			cluster                 entity.Cluster
			componentCapacityMapper = map[string]float64{}
		)
		err = db.Transaction(func(tx *gorm.DB) error {
			cluster, err = dbutil.GetCluster(tx, entity.ID(clusterID))
			if err != nil {
				return err
			}
			componentResources, err = clustermodel.GetClusterComponentsV2(cluster)
			if err != nil {
				return err
			}
			return nil
		})
		if err != nil {
			return &vec, err
		}

		for _, componentResource := range componentResources {
			if componentResource.Capacity.Compute != nil {
				componentCapacityMapper[componentResource.ComponentName.String()] = float64(componentResource.Capacity.Compute.CPU)
			}
		}
		// We can only get instance from thanos metrics.
		thanosNodeCPUSecondsRespResult, err := central_thanos.QueryNodeCPUUsedSeconds(ctx, clusterID)
		if err != nil || thanosNodeCPUSecondsRespResult.Status != "success" {
			return &vec, errors.Errorf("failed to query used cpu seconds, status: %s", thanosNodeCPUSecondsRespResult.Status)
		}

		for _, result := range thanosNodeCPUSecondsRespResult.Data.Result {
			m := result.LabelSet
			valueTuple := result.Value
			if len(valueTuple) != 2 {
				return &vec, errors.Errorf("failed to query value(%+v) of used cpu seconds", valueTuple)
			}

			if m.Component == component.PD.String() {
				continue
			}

			capacityValue := componentCapacityMapper[m.Component]
			vec.Metrics = append(vec.Metrics, &prom.VecMetricElement{
				Labels: []prom.MetricLabel{
					{Name: datasource.LabelInstance, Value: datasource.UnifyDBName(m.Instance)},
					{Name: datasource.LabelComponent, Value: datasource.UnifyDBName(m.Component)},
					{Name: datasource.LabelCluster, Value: cluster.Name},
				},
				Value: model.SampleValue(capacityValue),
			})
		}

		return &vec, nil
	}
}

func BuildNodeStorageCapacityQuerier(clusterID uint64) datasource.MetricQuerierFunc {
	return func(ctx context.Context) (prom.Metric, error) {
		vec := prom.VecMetric{
			Name:    "tidbcloud_node_storage_capacity_bytes",
			Help:    "The disk capacity bytes of TiKV/TiFlash nodes",
			Type:    "gauge",
			Metrics: []*prom.VecMetricElement{},
		}

		var (
			err                     error
			componentResources      []component.Resources
			cluster                 entity.Cluster
			componentCapacityMapper = map[string]uint64{}
		)
		err = db.Transaction(func(tx *gorm.DB) error {
			cluster, err = dbutil.GetCluster(tx, entity.ID(clusterID))
			if err != nil {
				return err
			}
			componentResources, err = clustermodel.GetClusterComponentsV2(cluster)
			if err != nil {
				return err
			}
			return nil
		})
		if err != nil {
			return &vec, err
		}

		for _, componentResource := range componentResources {
			if componentResource.Capacity.Disk != nil {
				componentCapacityMapper[componentResource.ComponentName.String()] = uint64(componentResource.Capacity.Disk.SizeGi) * 1024 * 1024 * 1024
			}
		}
		// We can only get instance from thanos metrics.
		thanosUsedNodeStorageRespResult, err := central_thanos.QueryNodeStorageUsedBytes(ctx, clusterID)
		if err != nil || thanosUsedNodeStorageRespResult.Status != "success" {
			return &vec, errors.Errorf("failed to query used node storage, status: %s", thanosUsedNodeStorageRespResult.Status)
		}

		for _, result := range thanosUsedNodeStorageRespResult.Data.Result {
			m := result.LabelSet
			valueTuple := result.Value
			if len(valueTuple) != 2 {
				return &vec, errors.Errorf("failed to query value(%+v) of used node storage", valueTuple)
			}

			if m.Component == component.PD.String() {
				continue
			}

			capacityValue := componentCapacityMapper[m.Component]
			vec.Metrics = append(vec.Metrics, &prom.VecMetricElement{
				Labels: []prom.MetricLabel{
					{Name: datasource.LabelInstance, Value: datasource.UnifyDBName(m.Instance)},
					{Name: datasource.LabelComponent, Value: datasource.UnifyDBName(m.Component)},
					{Name: datasource.LabelCluster, Value: cluster.Name},
				},
				Value: model.SampleValue(capacityValue),
			})
		}

		return &vec, nil
	}
}