package helper

import (
	relaycommon "github.com/QuantumNous/new-api/relay/common"
	"github.com/QuantumNous/new-api/relaykit/dto"
	"github.com/QuantumNous/new-api/setting/model_setting"
	"github.com/stretchr/testify/require"
	"testing"
)

func TestImageGenerationResolutionMapping(t *testing.T) {
	cfg := model_setting.GetGlobalSettings()
	oldGroup, oldModels := cfg.ImageGenerationGroup, cfg.ImageGenerationModels
	t.Cleanup(func() { cfg.ImageGenerationGroup, cfg.ImageGenerationModels = oldGroup, oldModels })
	cfg.ImageGenerationGroup, cfg.ImageGenerationModels = "gpt-images", []string{"gpt-image-2", "gpt-image-2-4k"}
	for _, tc := range []struct{ size, tier string }{{"", "1k"}, {"auto", "1k"}, {"1024x1536", "1k"}, {"1536x1024", "1k"}, {"2048x1152", "2k"}, {"2048x2048", "2k"}, {"3840x2160", "4k"}, {"2160x3840", "4k"}, {"4k", "4k"}} {
		info := &relaycommon.RelayInfo{OriginModelName: "gpt-image-2", UsingGroup: "gpt-images", ChannelMeta: &relaycommon.ChannelMeta{UpstreamModelName: "gpt-image-2"}}
		image := &dto.ImageRequest{Model: "gpt-image-2", Size: tc.size}
		require.NoError(t, mapImageGenerationResolution(info, image))
		require.Equal(t, "gpt-image-2-"+tc.tier, info.UpstreamModelName)
	}
	_, err := imageGenerationResolution("-1x1024")
	require.Error(t, err)
	_, err = imageGenerationResolution("garbage")
	require.Error(t, err)
	info := &relaycommon.RelayInfo{OriginModelName: "gpt-image-2", UsingGroup: "default", ChannelMeta: &relaycommon.ChannelMeta{UpstreamModelName: "gpt-image-2"}}
	require.NoError(t, mapImageGenerationResolution(info, &dto.ImageRequest{Model: "gpt-image-2", Size: "3840x2160"}))
	require.Equal(t, "gpt-image-2", info.UpstreamModelName)
}