package middleware

import (
	"context"

	"edge-infra.dev/pkg/lib/fog"

	"github.com/gin-contrib/requestid"
	"github.com/gin-gonic/gin"
	"github.com/go-logr/logr"
)

type correlationIDKey struct{}

const (
	CorrelationIDKey   = "X-Correlation-ID"
	correlationIDLabel = "correlationID"
)

// Sets the input logger with Operation ID and Correlation ID values in the server context, where
// Operation ID refers to each request that the component receives and Correlation ID refers to a
// request made by a user as it travels through the system.
//
// Expects the use of [pkg/edge-infra.dev/pkg/edge/api/middleware.SetRequestContext] middleware
// to set the Correlation ID value in the context. Please include WithCustomHeaderStrKey(CorrelationIDLabel)
// as an option in requestid.New().
//
// Expects the gin engine's ContextWithFallback field to be set to true to allow request
// context values to fallback to the parent context. By default it is false.
//
// The Correlation ID value is set from the incoming `X-Correlation-ID` header
// if present. It is possible to use the [GetCorrelationID] function to extract
// the incoming correlation ID from a context to be used to propagate to any
// further services used for a particular request
func SetLoggerInContext(log logr.Logger) gin.HandlerFunc {
	return func(c *gin.Context) {
		correlationID := requestid.Get(c)
		operationID := fog.OperationID(c.Request.Context())

		log := log.WithValues(correlationIDLabel, correlationID)
		log = log.WithValues(fog.OperationFields(operationID)...)

		ctxWithID := context.WithValue(c.Request.Context(), correlationIDKey{}, correlationID)
		c.Request = c.Request.Clone(fog.IntoContext(ctxWithID, log))

		c.Next()
	}
}

// Writes an initial log message before a request is processed and a final log after a request is finished (success OR fail).
// Should be set in a router after correlation and operation IDs are set, and before any middleware that could end the request.
func RequestBookendLogs() gin.HandlerFunc {
	return func(c *gin.Context) {
		// before request
		url := c.Request.URL.Path
		log := fog.FromContext(c, "endpoint", url)
		c.Request = c.Request.Clone(fog.IntoContext(c.Request.Context(), log))
		log.Info("Request received")

		c.Next()

		// after request
		status := c.Writer.Status()
		log = fog.FromContext(c, "statusCode", status)
		if status < 400 {
			log = log.WithValues("isSuccessful", true)
		} else {
			log = log.WithValues("isSuccessful", false)
		}
		log.Info("Request completed")
	}
}

func GetCorrelationID(ctx context.Context) string {
	correlationID := ctx.Value(correlationIDKey{})
	if correlationID == nil { // Avoid panic if correlationID has not been set
		return ""
	}
	return correlationID.(string)
}