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"
	"github.com/QuantumNous/new-api/setting/model_setting"
	"github.com/QuantumNous/new-api/setting/ratio_setting"
	"github.com/gin-gonic/gin"
	"github.com/stretchr/testify/require"
	"net/http/httptest"
	"testing"
)

func TestDedicatedImageGenerationPoolSelection(t *testing.T) {
	db := setupChannelSelectAutoGroupsTest(t)
	cfg := model_setting.GetGlobalSettings()
	oldGroup, oldModels := cfg.ImageGenerationGroup, cfg.ImageGenerationModels
	t.Cleanup(func() { cfg.ImageGenerationGroup, cfg.ImageGenerationModels = oldGroup, oldModels })
	cfg.ImageGenerationGroup, cfg.ImageGenerationModels = "example-images", []string{"gpt-image-2"}
	require.NoError(t, setting.UpdateUserUsableGroupsByJSONString(`{"default":"Default","vip":"VIP","example-images":"Images"}`))
	require.NoError(t, ratio_setting.UpdateGroupRatioByJSONString(`{"default":1,"vip":2,"example-images":1.25}`))
	createChannelSelectAutoGroupsChannel(t, db, 8001, "default", "gpt-image-2")
	createChannelSelectAutoGroupsChannel(t, db, 8002, "example-images", "gpt-image-2")
	model.InitChannelCache()
	ResetGroupAvailability()
	t.Cleanup(ResetGroupAvailability)
	for _, mode := range []string{"smart", "inherit", "custom"} {
		t.Run(mode, func(t *testing.T) {
			c, _ := gin.CreateTestContext(httptest.NewRecorder())
			c.Request = httptest.NewRequest("POST", "/v1/images/generations", nil)
			common.SetContextKey(c, constant.ContextKeyUserGroup, "default")
			c.Set(model.TokenRoutingContextKey, &model.TokenRoutingPolicy{Version: 1, Mode: mode, Groups: []model.TokenGroupEntry{{Group: "default", Priority: 0}}})
			entries := ResolveRequestPool(c, "gpt-image-2")
			require.Equal(t, []model.TokenGroupEntry{{Group: "example-images", Priority: 0}}, entries)
			channel, group, err := CacheGetRandomSatisfiedChannel(&RetryParam{Ctx: c, TokenGroup: "default", ModelName: "gpt-image-2", Retry: common.GetPointer(0)})
			require.NoError(t, err)
			require.NotNil(t, channel)
			require.Equal(t, 8002, channel.Id)
			require.Equal(t, "example-images", group)
		})
	}
	c, _ := gin.CreateTestContext(httptest.NewRecorder())
	c.Request = httptest.NewRequest("POST", "/v1/images/generations", nil)
	common.SetContextKey(c, constant.ContextKeyUserGroup, "default")
	defaults := []model.TokenGroupEntry{{Group: "default", Priority: 0}}
	c.Set(model.TokenRoutingContextKey, &model.TokenRoutingPolicy{Version: 1, Mode: "custom", Groups: defaults})
	require.Equal(t, defaults, ResolveRequestPool(c, "text-model"))
	c.Set(model.TokenRoutingContextKey, &model.TokenRoutingPolicy{Version: 1, Mode: "user"})
	require.Equal(t, defaults, ResolveRequestPool(c, "gpt-image-2"))
	c.Set(model.TokenRoutingContextKey, &model.TokenRoutingPolicy{Version: 1, Mode: "custom", Groups: defaults})
	require.NoError(t, setting.UpdateUserUsableGroupsByJSONString(`{"default":"Default","vip":"VIP"}`))
	require.Equal(t, defaults, ResolveRequestPool(c, "gpt-image-2"))
}