package middleware

import (
	"context"

	"github.com/gin-gonic/gin"
	"github.com/google/uuid"
)

const RequestIDHeader = "X-Request-ID"
const requestIDKey = "request_id"

type requestIDContextKey struct{}

func RequestID() gin.HandlerFunc {
	return func(c *gin.Context) {
		rid := c.GetHeader(RequestIDHeader)
		if rid == "" {
			rid = uuid.New().String()
		}
		c.Set(requestIDKey, rid)
		c.Request = c.Request.WithContext(context.WithValue(c.Request.Context(), requestIDContextKey{}, rid))
		c.Header(RequestIDHeader, rid)
		c.Next()
	}
}

func RequestIDFromGin(c *gin.Context) string {
	if rid, ok := c.Get(requestIDKey); ok {
		return rid.(string)
	}
	return ""
}

func RequestIDFromContext(ctx context.Context) string {
	if rid, ok := ctx.Value(requestIDContextKey{}).(string); ok {
		return rid
	}
	return ""
}