package controller

import (
	"github.com/QuantumNous/new-api/common"
	"github.com/QuantumNous/new-api/constant"
	"github.com/QuantumNous/new-api/middleware"
	"github.com/QuantumNous/new-api/model"
	"github.com/QuantumNous/new-api/setting"
	"github.com/QuantumNous/new-api/setting/ratio_setting"
	"github.com/gin-gonic/gin"
	"github.com/stretchr/testify/require"
	"net/http"
	"net/http/httptest"
	"testing"
)

func TestListModelsAutomaticPoolAllFormats(t *testing.T) {
	db := setupModelListControllerTestDB(t)
	withSelfUseModeDisabled(t)
	oldGroups, oldRatios, oldModelRatios := setting.UserUsableGroups2JSONString(), ratio_setting.GroupRatio2JSONString(), ratio_setting.ModelRatio2JSONString()
	oldMemory := common.MemoryCacheEnabled
	common.MemoryCacheEnabled = false
	t.Cleanup(func() {
		common.MemoryCacheEnabled = oldMemory
		require.NoError(t, setting.UpdateUserUsableGroupsByJSONString(oldGroups))
		require.NoError(t, ratio_setting.UpdateGroupRatioByJSONString(oldRatios))
		require.NoError(t, ratio_setting.UpdateModelRatioByJSONString(oldModelRatios))
	})
	require.NoError(t, setting.UpdateUserUsableGroupsByJSONString(`{"list-a":"","list-b":""}`))
	require.NoError(t, ratio_setting.UpdateGroupRatioByJSONString(`{"list-a":1,"list-b":1}`))
	require.NoError(t, ratio_setting.UpdateModelRatioByJSONString(`{"zz-list-a":1,"zz-list-b":1,"zz-shared":1,"zz-disabled":1}`))
	require.NoError(t, db.Create(&[]model.Channel{{Id: 5601, Type: constant.ChannelTypeOpenAI, Status: 1}, {Id: 5602, Type: constant.ChannelTypeOpenAI, Status: 1}, {Id: 5603, Status: 2}}).Error)
	require.NoError(t, db.Create(&[]model.Ability{
		{ChannelId: 5601, Group: "list-a", Model: "zz-list-a", Enabled: true},
		{ChannelId: 5602, Group: "list-b", Model: "zz-list-b", Enabled: true},
		{ChannelId: 5601, Group: "list-a", Model: "zz-shared", Enabled: true},
		{ChannelId: 5602, Group: "list-b", Model: "zz-shared", Enabled: true},
		{ChannelId: 5603, Group: "list-b", Model: "zz-disabled", Enabled: true},
		{ChannelId: 5602, Group: "list-b", Model: "zz-unpriced", Enabled: true},
	}).Error)
	for _, preference := range []string{"cost", "speed", "stable"} {
		for _, limited := range []bool{false, true} {
			for _, format := range []int{constant.ChannelTypeOpenAI, constant.ChannelTypeAnthropic, constant.ChannelTypeGemini} {
				policy, err := common.Marshal(model.TokenRoutingPolicy{Version: 1, Mode: "smart", AutoPreference: preference})
				require.NoError(t, err)
				token := &model.Token{Group: "auto", RoutingPolicy: string(policy), ModelLimitsEnabled: limited, ModelLimits: "zz-list-b,zz-disabled,zz-unpriced"}
				engine := gin.New()
				engine.GET("/v1/models", func(c *gin.Context) {
					common.SetContextKey(c, constant.ContextKeyUserGroup, "list-a")
					require.NoError(t, middleware.SetupContextForToken(c, token))
					ListModels(c, format)
				})
				res := httptest.NewRecorder()
				engine.ServeHTTP(res, httptest.NewRequest(http.MethodGet, "/v1/models", nil))
				require.Equal(t, 200, res.Code)
				var body struct {
					Data []struct {
						ID string `json:"id"`
					} `json:"data"`
					Models []struct {
						Name string `json:"name"`
					} `json:"models"`
				}
				require.NoError(t, common.Unmarshal(res.Body.Bytes(), &body))
				names := []string{}
				for _, item := range body.Data {
					names = append(names, item.ID)
				}
				for _, item := range body.Models {
					names = append(names, item.Name)
				}
				expected := []string{"zz-list-a", "zz-list-b", "zz-shared"}
				if limited {
					expected = []string{"zz-list-b"}
				}
				require.ElementsMatch(t, expected, names, "preference=%s limited=%v format=%d", preference, limited, format)
			}
		}
	}
}