package middleware

import (
	"github.com/labstack/echo/v4"

	"github.com/pingcap/test-infra/services/tcms/utils"
)

type RequestIDGenerator func() string

func RequestID(gen RequestIDGenerator) echo.MiddlewareFunc {
	return func(next echo.HandlerFunc) echo.HandlerFunc {
		return func(c echo.Context) error {
			// set reqID for Request if not present
			reqID := c.Request().Header.Get(echo.HeaderXRequestID)
			if reqID == "" {
				reqID = gen()
				c.Request().Header.Set(echo.HeaderXRequestID, reqID)
			}

			// set the same reqID for Response
			c.Response().Header().Set(echo.HeaderXRequestID, reqID)
			ctx := utils.NewContext(c)
			ctx.(utils.Context).SetLogReqID(reqID)
			return next(ctx)
		}
	}
}