package metrics

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

	"github.com/prometheus/client_golang/prometheus"
	"github.com/tidbcloud/tiinf/errs"
	"google.golang.org/grpc"
)

var (
	RpcCostHistogram *prometheus.HistogramVec
	RpcTotalCounter  *prometheus.CounterVec

	rpcRegisterOnce sync.Once
)

func rpcRegister(subsystem string) {
	rpcRegisterOnce.Do(func() {
		RpcTotalCounter = prometheus.NewCounterVec(prometheus.CounterOpts{
			Namespace: namespace,
			Subsystem: subsystem,
			Name:      "rpc_request_count",
			Help:      "The RPC request count.",
		}, []string{"service", "method", "code"})

		RpcCostHistogram = prometheus.NewHistogramVec(prometheus.HistogramOpts{
			Namespace: namespace,
			Subsystem: subsystem,
			Name:      "rpc_request_cost",
			Help:      "The RPC request latencies in seconds",
			Buckets:   prometheus.ExponentialBuckets(0.001, 2, 20), // 1ms ~ 524s
		}, []string{"service", "method"})

		prometheus.MustRegister(
			RpcTotalCounter,
			RpcCostHistogram,
		)
	})
}

func UnaryServerInterceptor() grpc.UnaryServerInterceptor {
	return func(ctx context.Context, req interface{}, info *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (interface{}, error) {
		var (
			serviceName string
			methodName  string
			code        string
		)
		serviceName, methodName = splitMethodName(info.FullMethod)
		startTime := time.Now()
		resp, err := handler(ctx, req)

		if baseResp := errs.GetBaseResp(resp); baseResp != nil {
			code = strconv.FormatInt(baseResp.ErrCode, 10)
		} else if err != nil {
			code = "error"
		} else {
			code = "0"
		}
		RpcTotalCounter.WithLabelValues(serviceName, methodName, code).Inc()
		RpcCostHistogram.WithLabelValues(serviceName, methodName).Observe(time.Since(startTime).Seconds())
		return resp, err
	}
}

func splitMethodName(fullMethodName string) (string, string) {
	fullMethodName = strings.TrimPrefix(fullMethodName, "/") // remove leading slash
	if i := strings.Index(fullMethodName, "/"); i >= 0 {
		return fullMethodName[:i], fullMethodName[i+1:]
	}
	return "unknown", "unknown"
}