package controller

import (
	"net/http"
	"net/http/httptest"
	"testing"

	"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/ratio_setting"
	"github.com/gin-gonic/gin"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
)

func TestSelectableGroupsRequireEnabledChannel(t *testing.T) {
	originalUsableGroups := setting.UserUsableGroups2JSONString()
	originalRatios := ratio_setting.GroupRatio2JSONString()
	require.NoError(t, setting.UpdateUserUsableGroupsByJSONString(`{"auto":"Auto","active":"Active","default":"Default","paused":"Paused","no-channel":"No channel"}`))
	require.NoError(t, ratio_setting.UpdateGroupRatioByJSONString(`{"active":0.2,"default":1,"paused":0.3,"no-channel":0.4}`))
	t.Cleanup(func() {
		require.NoError(t, setting.UpdateUserUsableGroupsByJSONString(originalUsableGroups))
		require.NoError(t, ratio_setting.UpdateGroupRatioByJSONString(originalRatios))
		model.InvalidatePricingCache()
	})

	db := setupModelListControllerTestDB(t)
	require.NoError(t, db.Create(&model.User{
		Id:       7001,
		Username: "group-availability-user",
		Password: "password",
		Group:    "default",
		Status:   common.UserStatusEnabled,
	}).Error)
	require.NoError(t, db.Create(&[]model.Channel{
		{Id: 11, Type: constant.ChannelTypeOpenAI, Key: "active-key", Name: "active", Status: common.ChannelStatusEnabled},
		{Id: 12, Type: constant.ChannelTypeOpenAI, Key: "default-key", Name: "default", Status: common.ChannelStatusManuallyDisabled},
		{Id: 13, Type: constant.ChannelTypeOpenAI, Key: "paused-key", Name: "paused", Status: common.ChannelStatusAutoDisabled},
	}).Error)
	require.NoError(t, db.Create(&[]model.Ability{
		{Group: "active", Model: "gpt-4o", ChannelId: 11, Enabled: true},
		{Group: "default", Model: "gpt-4o", ChannelId: 12, Enabled: true},
		{Group: "paused", Model: "gpt-4o", ChannelId: 13, Enabled: true},
	}).Error)
	model.InvalidatePricingCache()

	userRecorder := httptest.NewRecorder()
	userContext, _ := gin.CreateTestContext(userRecorder)
	userContext.Request = httptest.NewRequest(http.MethodGet, "/api/user/groups", nil)
	userContext.Set("id", 7001)
	GetUserGroups(userContext)

	var userResponse struct {
		Success bool                              `json:"success"`
		Data    map[string]map[string]interface{} `json:"data"`
	}
	require.NoError(t, common.Unmarshal(userRecorder.Body.Bytes(), &userResponse))
	require.True(t, userResponse.Success)
	assert.ElementsMatch(t, []string{"active", "auto"}, mapKeys(userResponse.Data))

	pricingRecorder := httptest.NewRecorder()
	pricingContext, _ := gin.CreateTestContext(pricingRecorder)
	pricingContext.Request = httptest.NewRequest(http.MethodGet, "/api/pricing", nil)
	GetPricing(pricingContext)

	var pricingResponse struct {
		Success     bool               `json:"success"`
		GroupRatio  map[string]float64 `json:"group_ratio"`
		UsableGroup map[string]string  `json:"usable_group"`
	}
	require.NoError(t, common.Unmarshal(pricingRecorder.Body.Bytes(), &pricingResponse))
	require.True(t, pricingResponse.Success)
	assert.Equal(t, map[string]float64{"active": 0.2}, pricingResponse.GroupRatio)
	assert.Equal(t, map[string]string{"active": "Active", "auto": "Auto"}, pricingResponse.UsableGroup)
}

func mapKeys[V any](values map[string]V) []string {
	keys := make([]string, 0, len(values))
	for key := range values {
		keys = append(keys, key)
	}
	return keys
}