package middleware

import (
	"fmt"
	"strings"

	"github.com/QuantumNous/new-api/common"
	"github.com/QuantumNous/new-api/constant"
	"github.com/gin-gonic/gin"
)

const (
	museHeaderUserID  = "X-Muse-User-Id"
	museHeaderWorkID  = "X-Muse-Work-Id"
	museHeaderRequest = "X-Request-Id"
	museHeaderTraceID = "X-Trace-Id"
	museHeaderScene   = "X-Muse-Scene"
	jsonContentType   = "application/json"
)

type museMetadataEnvelope struct {
	Metadata map[string]any `json:"metadata"`
}

func MuseRequestContext() gin.HandlerFunc {
	return func(c *gin.Context) {
		metadata := extractMuseMetadata(c)

		museUserID := firstNonBlank(
			c.GetHeader(museHeaderUserID),
			stringifyMetadata(metadata["muse_user_id"]),
		)
		museWorkID := firstNonBlank(
			c.GetHeader(museHeaderWorkID),
			stringifyMetadata(metadata["work_id"]),
			stringifyMetadata(metadata["muse_work_id"]),
		)
		museScene := firstNonBlank(
			c.GetHeader(museHeaderScene),
			stringifyMetadata(metadata["scene"]),
		)
		requestID := firstNonBlank(
			c.GetHeader(museHeaderRequest),
			stringifyMetadata(metadata["request_id"]),
		)
		traceID := firstNonBlank(
			c.GetHeader(museHeaderTraceID),
			stringifyMetadata(metadata["trace_id"]),
		)

		if museUserID != "" {
			common.SetContextKey(c, constant.ContextKeyMuseUserID, museUserID)
		}
		if museWorkID != "" {
			common.SetContextKey(c, constant.ContextKeyMuseWorkID, museWorkID)
		}
		if requestID != "" {
			common.SetContextKey(c, constant.ContextKeyMuseRequestID, requestID)
		}
		if museScene != "" {
			common.SetContextKey(c, constant.ContextKeyMuseScene, museScene)
		}
		if traceID != "" {
			common.SetContextKey(c, constant.ContextKeyMuseTraceID, traceID)
		}
		c.Next()
	}
}

func extractMuseMetadata(c *gin.Context) map[string]any {
	if c == nil || c.Request == nil {
		return nil
	}
	if !strings.Contains(strings.ToLower(c.GetHeader("Content-Type")), jsonContentType) {
		return nil
	}
	if !needsMuseMetadataRead(c) {
		return nil
	}
	var envelope museMetadataEnvelope
	if err := common.UnmarshalBodyReusable(c, &envelope); err != nil {
		return nil
	}
	return envelope.Metadata
}

func needsMuseMetadataRead(c *gin.Context) bool {
	return c.GetHeader(museHeaderUserID) == "" ||
		c.GetHeader(museHeaderWorkID) == "" ||
		c.GetHeader(museHeaderRequest) == "" ||
		c.GetHeader(museHeaderTraceID) == "" ||
		c.GetHeader(museHeaderScene) == ""
}

func firstNonBlank(values ...string) string {
	for _, value := range values {
		if strings.TrimSpace(value) != "" {
			return strings.TrimSpace(value)
		}
	}
	return ""
}

func stringifyMetadata(value any) string {
	switch v := value.(type) {
	case nil:
		return ""
	case string:
		return strings.TrimSpace(v)
	default:
		return strings.TrimSpace(fmt.Sprintf("%v", v))
	}
}