package middleware

import (
	"github.com/gin-gonic/gin"
	"github.com/pingcap/dbaas/billing2/utils/header"
	"net/http"
	"strings"
)

const (
	PressureTestHeaderKey = "X-PRESSURE-TEST"
	FunctionTestHeaderKey = "X-FUNCTION-TEST"
	TraceInfoHeaderKey    = "B3"
)

func parseFieldsInfo(c *gin.Context) {
	pressureTestTag, ok := c.Request.Header[http.CanonicalHeaderKey(PressureTestHeaderKey)]
	// header field of pressure test exists
	if ok && len(pressureTestTag) != 0 {
		header.SetPressureTestValue(pressureTestTag[0])
	}
	functionTestTag, ok := c.Request.Header[http.CanonicalHeaderKey(FunctionTestHeaderKey)]
	// header field of function test exists
	if ok && len(functionTestTag) != 0 {
		header.SetFunctionTestValue(functionTestTag[0])
	}
	// get traceID
	B3, ok := c.Request.Header[http.CanonicalHeaderKey(TraceInfoHeaderKey)]
	if ok && len(B3) != 0 {
		traceArray := strings.Split(B3[0], "-")
		if len(traceArray) != 0 {
			header.SetTraceID(traceArray[0])
		}
	}
}

func HeaderMiddleware(c *gin.Context) {
	parseFieldsInfo(c)
	c.Next()
}