package trace

import (
	"context"
	"fmt"
	"gitea.loveuer.com/yizhisec/packages/opt"
	"github.com/gin-gonic/gin"
	"github.com/google/uuid"
	"time"
)

func New() gin.HandlerFunc {
	return func(c *gin.Context) {
		trace := c.GetHeader(opt.TraceKey)
		if trace == "" {
			if u, err := uuid.NewV7(); err == nil {
				trace = u.String()
			} else {
				trace = fmt.Sprintf("%d", time.Now().UnixNano())
			}
		}

		ctx := context.WithValue(c.Request.Context(), opt.TraceKey, trace)
		c.Request = c.Request.WithContext(ctx)

		c.Writer.Header().Set(opt.TraceKey, trace)

		c.Next()
	}
}