package middleware

import (
	"net/http"

	"github.com/labstack/echo/v4"

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

// HeaderKey is the specified key for the Server-Timing header.
const HeaderKey = "Server-Timing"

func ServerTiming(next echo.HandlerFunc) echo.HandlerFunc {
	return func(c echo.Context) error {
		h := timing.NewHeader()
		r := c.Request()
		c.SetRequest(r.WithContext(timing.NewContext(r.Context(), h)))

		totalTime := timing.Total(c)
		c.Response().Before(func() {
			// Record total timing just before the response is written.
			totalTime.End()
			writeHeader(c.Response().Header(), h)
		})
		return next(c)
	}
}

func writeHeader(headers http.Header, h *timing.Header) {
	h.Lock()
	defer h.Unlock()
	// If there are no metrics set, do nothing
	if h.Empty() {
		return
	}
	headers.Set(HeaderKey, h.String())
}