package helper

import (
	"fmt"
	relaycommon "github.com/QuantumNous/new-api/relay/common"
	"github.com/QuantumNous/new-api/relaykit/dto"
	"github.com/QuantumNous/new-api/setting/model_setting"
	"slices"
	"strconv"
	"strings"
)

func imageGenerationResolution(size string) (string, error) {
	size = strings.ToLower(strings.TrimSpace(size))
	if size == "" || size == "auto" || size == "1k" {
		return "1k", nil
	}
	if size == "2k" || size == "4k" {
		return size, nil
	}
	dimensions := strings.Split(size, "x")
	if len(dimensions) != 2 {
		return "", fmt.Errorf("size must be auto, 1k, 2k, 4k or WIDTHxHEIGHT")
	}
	width, errWidth := strconv.Atoi(dimensions[0])
	height, errHeight := strconv.Atoi(dimensions[1])
	if errWidth != nil || errHeight != nil || width <= 0 || height <= 0 {
		return "", fmt.Errorf("invalid image size")
	}
	short := min(width, height)
	if short <= 1024 {
		return "1k", nil
	}
	if short <= 2048 {
		return "2k", nil
	}
	return "4k", nil
}

func mapImageGenerationResolution(info *relaycommon.RelayInfo, request dto.Request) error {
	cfg := model_setting.GetGlobalSettings()
	image, ok := request.(*dto.ImageRequest)
	if !ok || cfg.ImageGenerationGroup == "" || info.UsingGroup != cfg.ImageGenerationGroup || !slices.Contains(cfg.ImageGenerationModels, info.OriginModelName) {
		return nil
	}
	name := info.OriginModelName
	for _, suffix := range []string{"-1k", "-2k", "-4k"} {
		if strings.HasSuffix(name, suffix) {
			return nil
		}
	}
	tier, err := imageGenerationResolution(image.Size)
	if err != nil {
		return err
	}
	info.UpstreamModelName = name + "-" + tier
	info.IsModelMapped = true
	return nil
}