package metric

import (
	"context"
	"fmt"
	"github.com/guo-shaoge/tiflash-autoscale/pkg/util/log/logger"
	promapi "github.com/prometheus/client_golang/api"
	"github.com/prometheus/client_golang/api/prometheus/v1"
	"github.com/prometheus/common/model"
	"go.uber.org/zap"
	"time"
)

// reference:
//  1. example code: https://github.com/prometheus/client_golang/blob/main/api/prometheus/v1/example_test.go
//  2. golang doc: https://pkg.go.dev/github.com/prometheus/client_golang/api/prometheus/v1
type Collector struct {
	ctx                 context.Context
	promCli             promapi.Client
	promQueryTimeoutSec uint
	logger              *zap.Logger
}

var (
	// cpuPromQL     = "container_cpu_usage_seconds_total{job=\"kube_sd\", pod=~\"readnode.+\",container=\"supervisor\"}[1m]"
	cpuPromQL = "tiflash_proxy_process_cpu_seconds_total{job=\"kube_sd_tiflash_proc\",pod=~\"readnode.+\"}[1m]"
	oomPromQL = "sum by(pod) (max_over_time(tiflash_memory_exceed_quota_count{job=\"kube_sd_tiflash_proc\",metrics_topic=\"tiflash\", pod!=\"\"}[30s]))"
	// taskCntPromQL = "sum by(pod) (sum_over_time(tiflash_coprocessor_handling_request_count{job=\"kube_sd_tiflash_proc\",metrics_topic=\"tiflash\", pod!=\"\"}[30s]))"
	// taskCntPromQL = "sum by(pod,resource_group) (sum_over_time(tiflash_resource_group{job=\"kube_sd_tiflash_proc\",metrics_topic=\"tiflash\", pod!=\"\", type=~\"handling_mpp_task_run|handling_mpp_task_dispatch|handling_mpp_task_establish|handling_mpp_task_cancel\"}[30s]))"
	taskCntPromQL = "sum by(pod,resource_group) (tiflash_request_count_per_resource_group{job=\"kube_sd_tiflash_proc\",metrics_topic=\"tiflash\", pod!=\"\", type=~\"mpp_task_run|mpp_task_dispatch|mpp_task_establish|mpp_task_cancel\"})"
	ruPromQL      = "sum_over_time(tiflash_compute_request_unit{job=\"kube_sd_tiflash_proc\",metrics_topic=\"tiflash\", pod!=\"\",type=\"mpp\"}[30s])"
)

func NewCollector(ctx context.Context, promAddr string, promQueryTimeoutSec uint, parentLogger *zap.Logger) (*Collector, error) {
	promCli, err := promapi.NewClient(promapi.Config{
		Address: promAddr,
	})
	if err != nil {
		return nil, err
	}
	res := &Collector{
		ctx:                 ctx,
		promCli:             promCli,
		promQueryTimeoutSec: promQueryTimeoutSec,
		logger:              parentLogger.With(zap.String(logger.LogIDMod, "metricCollector")),
	}
	return res, nil
}

func (c *Collector) doQuery(query string) (model.Value, error) {
	v1api1 := v1.NewAPI(c.promCli)
	ctx, cancel := context.WithTimeout(c.ctx, time.Duration(c.promQueryTimeoutSec)*time.Second)
	defer cancel()

	queryRes, warning, err := v1api1.Query(ctx, query, time.Now(), v1.WithTimeout(5*time.Second))
	if err != nil {
		return nil, err
	}

	if len(warning) != 0 {
		c.logger.Warn("got warning, ignore", zap.Any("warnings", warning))
	}
	return queryRes, err
}

// pod-i: [cpu-1, cpu-2,...]
// pod-j: [cpu-1, cpu-2,...]
// "queryRes":
//
//	  "tiflash_proxy_process_cpu_seconds_total{
//			addr="10.244.0.210:8234",
//			app="readnode",
//			app_kubernetes_io_component="tiflash-autoscale",
//			apps_kruise_io_cloneset_instance_id="pksg7",
//			controller_revision_hash="readnode-7c455c9768",
//			free_tier_tidbcloud_com_mtls="false",
//			instance="10.244.0.210:8234",
//			job="kube_sd_tiflash_proc",
//			lifecycle_apps_kruise_io_state="Normal",
//			metrics_topic="tiflash",
//			mpath="/metrics",
//			pod="readnode-pksg7",
//			pod_ip="10.244.0.210",
//			pod_template_hash="7c455c9768",
//			prom_io_path="/metrics",
//	     sidecar_istio_io_inject="false"} =>\n518
//
// @[1725455210.628]
// 519 @[1725455220.628]
// 532 @[1725455230.628]
// 546 @[1725455240.628]
// 552 @[1725455250.628]
// 553 @[1725455260.628]
//
//	tiflash_proxy_process_cpu_seconds_total{
//			addr="10.244.0.213:8234",
//			app="readnode",
//			app_kubernetes_io_component="tiflash-autoscale",
//			apps_kruise_io_cloneset_instance_id="cv4wd",
//			controller_revision_hash="readnode-7c455c9768",
//			free_tier_tidbcloud_com_mtls="false",
//			instance="10.244.0.213:8234",
//			job="kube_sd_tiflash_proc",
//			lifecycle_apps_kruise_io_state="Normal",
//			metrics_topic="tiflash",
//			mpath="/metrics",
//			pod="readnode-cv4wd",
//			pod_ip="10.244.0.213",
//			pod_template_hash="7c455c9768",
//			prom_io_path="/metrics", sidecar_istio_io_inject="false"} =>4 @[1725455262.104]"}
func (c *Collector) QueryCPU(podsWindow PodsCPUMetricWindow) (err error) {
	queryRes, err := c.doQuery(cpuPromQL)
	defer func() {
		var queryResStr string
		if queryRes != nil {
			queryResStr = queryRes.String()
		}
		c.logger.Info("cpu metric window info", zap.Any("window", podsWindow.ToString()),
			zap.Any("queryRes", queryResStr), zap.Error(err))
	}()
	if err != nil {
		return err
	}

	matrix, ok := queryRes.(model.Matrix)
	if !ok {
		err = fmt.Errorf("metric is not Matrix, type: %v", queryRes.Type())
		return err
	}

	// stores one point for each pods.
	podsCPUPointMap := make(map[string]*CPUPoint)
	for _, podSample := range matrix {
		podName := string(podSample.Metric["pod"])
		podSampleLen := len(podSample.Values)
		if podSampleLen < 2 {
			c.logger.Info("pod metric point not enough, skip", zap.String("pod", podName))
			continue
		}

		last := podSample.Values[podSampleLen-1]
		prevLast := podSample.Values[podSampleLen-2]
		dura := float64(last.Timestamp.Unix() - prevLast.Timestamp.Unix())
		rate := float64(last.Value-prevLast.Value) / dura
		if rate < 0.0 {
			c.logger.Info("got negative cpu info, ignore",
				zap.Any("rate", rate),
				zap.Any("prevLast", prevLast),
				zap.Any("last", last),
				zap.Any("dura", dura),
				zap.Any("calc", fmt.Sprintf("%v-%v=%v", last.Value, prevLast.Value, float64(last.Value-prevLast.Value))))
			continue
		}

		newPoint := CPUPoint{
			basePoint: basePoint{
				ts: last.Timestamp.Time(),
			},
			cpu: rate,
		}

		if _, ok := podsCPUPointMap[podName]; ok {
			err = fmt.Errorf("duplicated cpu point for pod(%s)", podName)
			return err
		} else {
			podsCPUPointMap[podName] = &newPoint
		}
	}

	// merge one metric point into metric window for each pod.
	for podName, point := range podsCPUPointMap {
		window, ok := podsWindow[podName]
		if ok {
			if len(window) > 0 {
				l := window[len(window)-1]
				if l.GetTs().After(point.GetTs()) {
					c.logger.Info("new query cpu metric point should be after recorded one",
						zap.Any("new ts", point.GetTs()), zap.Any("recorded ts", l.GetTs()))
					continue
				}
			}
			window = append(window, point)
			podsWindow[podName] = window
		} else {
			podsWindow[podName] = CPUMetricWindow{point}
		}
	}
	return nil
}

func (c *Collector) QueryOOM(podsWindow PodsOOMMetricWindow) (err error) {
	queryRes, err := c.doQuery(oomPromQL)
	defer func() {
		var queryResStr string
		if queryRes != nil {
			queryResStr = queryRes.String()
		}
		c.logger.Info("oom metric window info", zap.Any("window", podsWindow.ToString()),
			zap.Any("queryRes", queryResStr), zap.Error(err))
	}()
	if err != nil {
		return err
	}

	vector, ok := queryRes.(model.Vector)
	if !ok {
		err = fmt.Errorf("metric is not Vector, type: %v", queryRes.Type())
		return err
	}

	// one metric point for all pods.
	podsOOMPointMap := make(map[string]*OOMPoint)
	for _, sample := range vector {
		podName := string(sample.Metric["pod"])
		newPoint := OOMPoint{
			basePoint: basePoint{
				ts: sample.Timestamp.Time(),
			},
			oom: uint(float64(sample.Value)),
		}

		_, ok := podsOOMPointMap[podName]
		if ok {
			err = fmt.Errorf("duplicated oom point for pod(%s)", podName)
			return err
		} else {
			podsOOMPointMap[podName] = &newPoint
		}
	}

	// merge one metric point into metric window for each pod.
	for podName, point := range podsOOMPointMap {
		window, ok := podsWindow[podName]
		if ok {
			if len(window) > 0 {
				l := window[len(window)-1]
				if l.GetTs().After(point.GetTs()) {
					c.logger.Info("new query cpu metric point should be after recorded one",
						zap.Any("new ts", point.GetTs()), zap.Any("recorded ts", l.GetTs()))
					continue
				}
			}
			window = append(window, point)
			podsWindow[podName] = window
		} else {
			podsWindow[podName] = OOMMetricWindow{point}
		}
	}
	return nil
}

// func (c *Collector) QueryTaskCnt(podsWindow PodsTaskCntMetricWindow) (err error) {
// 	queryRes, err := c.doQuery(taskCntPromQL)
// 	defer func() {
// 		var queryResStr string
// 		if queryRes != nil {
// 			queryResStr = queryRes.String()
// 		}
// 		c.logger.Info("taskCnt metric window info", zap.Any("window", podsWindow.ToString()),
// 			zap.Any("queryRes", queryResStr), zap.Error(err))
// 	}()
// 	if err != nil {
// 		return err
// 	}
//
// 	vector, ok := queryRes.(model.Vector)
// 	if !ok {
// 		err = fmt.Errorf("metric is not Vector, type: %v", queryRes.Type())
// 		return err
// 	}
//
// 	// one metric point for all pods.
// 	podsTaskCntPointMap := make(map[string]*TaskCntPoint)
// 	for _, sample := range vector {
// 		podName := string(sample.Metric["pod"])
// 		newPoint := TaskCntPoint{
// 			basePoint: basePoint{
// 				ts: sample.Timestamp.Time(),
// 			},
// 			taskCnt: uint(float64(sample.Value)),
// 		}
//
// 		_, ok := podsTaskCntPointMap[podName]
// 		if ok {
// 			err = fmt.Errorf("duplicated taskCnt point for pod(%s)", podName)
// 			return err
// 		} else {
// 			podsTaskCntPointMap[podName] = &newPoint
// 		}
// 	}
//
// 	// merge one metric point into metric window for each pod.
// 	for podName, point := range podsTaskCntPointMap {
// 		window, ok := podsWindow[podName]
// 		if ok {
// 			if len(window) > 0 {
// 				l := window[len(window)-1]
// 				if l.GetTs().After(point.GetTs()) {
// 					c.logger.Info("new query cpu metric point should be after recorded one",
// 						zap.Any("new ts", point.GetTs()), zap.Any("recorded ts", l.GetTs()))
// 					continue
// 				}
// 			}
// 			window = append(window, point)
// 			podsWindow[podName] = window
// 		} else {
// 			podsWindow[podName] = TaskCntMetricWindow{point}
// 		}
// 	}
// 	return nil
// }

func (c *Collector) QueryTaskCnt(podsWindow PodsTaskCntMetricWindow) (err error) {
	queryRes, err := c.doQuery(taskCntPromQL)
	defer func() {
		var queryResStr string
		if queryRes != nil {
			queryResStr = queryRes.String()
		}
		c.logger.Info("taskCnt metric window info", zap.Any("window", podsWindow.ToString()),
			zap.Any("queryRes", queryResStr), zap.Error(err))
	}()
	if err != nil {
		return err
	}

	vector, ok := queryRes.(model.Vector)
	if !ok {
		err = fmt.Errorf("metric is not Vector, type: %v", queryRes.Type())
		return err
	}

	// one metric point for all pods.
	podsTaskCntPointMap := make(map[string]*TaskCntPoint)
	for _, sample := range vector {
		podName := string(sample.Metric["pod"])
		rgName := ResourceGroupNameType(sample.Metric["resource_group"])
		tenantTaskCnt := uint(float64(sample.Value))
		taskCntPoint, ok := podsTaskCntPointMap[podName]
		if ok {
			if _, ok := taskCntPoint.tenantTaskCntMap[rgName]; ok {
				err = fmt.Errorf("duplicated taskCnt point for pod(%s) and tenant(%s)", podName, rgName)
				return err
			}
			taskCntPoint.totalTaskCnt += tenantTaskCnt
			taskCntPoint.tenantTaskCntMap[rgName] = tenantTaskCnt
		} else {
			podsTaskCntPointMap[podName] = &TaskCntPoint{
				basePoint: basePoint{
					ts: sample.Timestamp.Time(),
				},
				totalTaskCnt: tenantTaskCnt,
				tenantTaskCntMap: map[ResourceGroupNameType]uint{
					rgName: tenantTaskCnt,
				},
			}
		}
	}

	// merge that one metric point into metric window for each pod.
	for podName, point := range podsTaskCntPointMap {
		window, ok := podsWindow[podName]
		if ok {
			if len(window) > 0 {
				l := window[len(window)-1]
				if l.GetTs().After(point.GetTs()) {
					c.logger.Info("new query cpu metric point should be after recorded one",
						zap.Any("new ts", point.GetTs()), zap.Any("recorded ts", l.GetTs()))
					continue
				}
			}
			window = append(window, point)
			podsWindow[podName] = window
		} else {
			podsWindow[podName] = TaskCntMetricWindow{point}
		}
	}
	return nil
}

func (c *Collector) QueryRU(podsWindow PodsRUMetricWindow) (err error) {
	queryRes, err := c.doQuery(ruPromQL)
	defer func() {
		var queryResStr string
		if queryRes != nil {
			queryResStr = queryRes.String()
		}
		c.logger.Info("ru metric window info", zap.Any("window", podsWindow.ToString()),
			zap.Any("queryRes", queryResStr), zap.Error(err))
	}()
	if err != nil {
		return err
	}

	vector, ok := queryRes.(model.Vector)
	if !ok {
		err = fmt.Errorf("metric is not Vector, type: %v", queryRes.Type())
		return err
	}

	// one metric point for all pods.
	podsRUPointMap := make(map[string]*RUPoint)
	for _, sample := range vector {
		podName := string(sample.Metric["pod"])
		rgName := ResourceGroupNameType(sample.Metric["resource_group"])
		tenantRU := float64(sample.Value)
		ruPoint, ok := podsRUPointMap[podName]
		if ok {
			if _, ok := ruPoint.tenantRUMap[rgName]; ok {
				err = fmt.Errorf("duplicated ru point for pod(%s) and tenant(%s)", podName, rgName)
				return err
			}
			ruPoint.totalRU += tenantRU
			ruPoint.tenantRUMap[rgName] = tenantRU
		} else {
			podsRUPointMap[podName] = &RUPoint{
				basePoint: basePoint{
					ts: sample.Timestamp.Time(),
				},
				totalRU: tenantRU,
				tenantRUMap: map[ResourceGroupNameType]float64{
					rgName: tenantRU,
				},
			}
		}
	}

	// merge that one metric point into metric window for each pod.
	for podName, point := range podsRUPointMap {
		window, ok := podsWindow[podName]
		if ok {
			if len(window) > 0 {
				l := window[len(window)-1]
				if l.GetTs().After(point.GetTs()) {
					c.logger.Info("new query cpu metric point should be after recorded one",
						zap.Any("new ts", point.GetTs()), zap.Any("recorded ts", l.GetTs()))
					continue
				}
			}
			window = append(window, point)
			podsWindow[podName] = window
		} else {
			podsWindow[podName] = RUMetricWindow{point}
		}
	}
	return nil
}