package service

import (
	"github.com/QuantumNous/new-api/common"
	"github.com/QuantumNous/new-api/constant"
	"github.com/QuantumNous/new-api/model"
	"github.com/QuantumNous/new-api/setting/model_setting"
	"github.com/gin-gonic/gin"
	"slices"
)

func imageGenerationPoolGroup(c *gin.Context) string {
	settings := model_setting.GetGlobalSettings()
	group := settings.ImageGenerationGroup
	if group == "" || !IsUserSelectableGroup(common.GetContextKeyString(c, constant.ContextKeyUserGroup), group) {
		return ""
	}
	policy := RequestRoutingPolicy(c)
	if policy != nil {
		if policy.Mode == "user" {
			return ""
		}
		return group
	}
	if common.GetContextKeyString(c, constant.ContextKeyTokenGroup) == "auto" || common.GetContextKeyBool(c, constant.ContextKeyTokenCrossGroupRetry) {
		return group
	}
	return ""
}

func imageGenerationRequestGroup(c *gin.Context, modelName string) string {
	if !slices.Contains(model_setting.GetGlobalSettings().ImageGenerationModels, modelName) {
		return ""
	}
	return imageGenerationPoolGroup(c)
}

// The dedicated route is shared by selection, affinity and the routing preview.
func WithImageGenerationGroup(c *gin.Context, modelName string, entries []model.TokenGroupEntry) []model.TokenGroupEntry {
	group := imageGenerationPoolGroup(c)
	if group == "" {
		return entries
	}
	if modelName != "" {
		if imageGenerationRequestGroup(c, modelName) != "" {
			return []model.TokenGroupEntry{{Group: group, Priority: 0}}
		}
		return entries
	}
	priority := 0
	for _, entry := range entries {
		if entry.Group == group {
			return entries
		}
		priority = max(priority, entry.Priority+1)
	}
	return append(slices.Clone(entries), model.TokenGroupEntry{Group: group, Priority: priority})
}