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/operation_setting"
	"github.com/gin-gonic/gin"
	"github.com/stretchr/testify/require"
	"net/http/httptest"
	"testing"
)

func TestTokenRoutingPoolPreferenceExclusionAndAffinityIsolation(t *testing.T) {
	db := setupChannelSelectAutoGroupsTest(t)
	const modelName = "pool-model"
	createChannelSelectAutoGroupsChannel(t, db, 6101, "default", modelName)
	createChannelSelectAutoGroupsChannel(t, db, 6102, "vip", modelName)
	model.InitChannelCache()
	ResetGroupAvailability()
	t.Cleanup(ResetGroupAvailability)
	config := operation_setting.GetChannelAffinitySetting()
	previous := *config
	*config = operation_setting.ChannelAffinitySetting{Enabled: true, DefaultTTLSeconds: 3600, SessionMode: "prefer", Rules: []operation_setting.ChannelAffinityRule{{Name: "pool-test", ModelRegex: []string{".*"}, KeySources: []operation_setting.ChannelAffinityKeySource{{Type: "request_header", Key: "Session-Id"}}, SessionMode: "prefer"}}}
	t.Cleanup(func() { *config = previous; _ = poolBindings().Purge() })
	ctx := func(token int, session string) *gin.Context {
		c := setupGroupPoolContext(t)
		c.Request = httptest.NewRequest("POST", "/v1/responses", nil)
		c.Request.Header.Set("Session-Id", session)
		c.Set("id", 100)
		c.Set("token_id", token)
		c.Set(model.TokenRoutingContextKey, &model.TokenRoutingPolicy{Version: 1, Mode: "custom", Groups: []model.TokenGroupEntry{{Group: "default", Priority: 0}, {Group: "vip", Priority: 1}}, ModelRules: []model.ModelRoutingRule{{Model: modelName, Groups: []model.TokenGroupEntry{{Group: "vip", Priority: 0}}}}})
		return c
	}
	c := ctx(100, "session-a")
	param := &RetryParam{Ctx: c, TokenGroup: "default", ModelName: modelName, Retry: common.GetPointer(0)}
	first, group, err := SelectChannelForRequest(c, modelName, param)
	require.Nil(t, err)
	require.NotNil(t, first)
	require.Equal(t, 6102, first.Id)
	require.Equal(t, "vip", group)
	AppendUsedChannel(c, first.Id)
	second, group, err2 := CacheGetRandomSatisfiedChannel(param)
	require.NoError(t, err2)
	require.NotNil(t, second)
	require.Equal(t, 6101, second.Id)
	c.Set("channel_id", second.Id)
	common.SetContextKey(c, constant.ContextKeyAutoGroup, group)
	RequestPolicy(c).Successful = true
	RecordPoolAffinity(c)
	next := ctx(100, "session-a")
	bound, boundGroup := PreferredPoolChannel(next, modelName)
	require.NotNil(t, bound)
	require.Equal(t, 6101, bound.Id)
	require.Equal(t, "default", boundGroup)
	isolated, _ := PreferredPoolChannel(ctx(101, "session-a"), modelName)
	require.Nil(t, isolated)
	isolated, _ = PreferredPoolChannel(ctx(100, "session-b"), modelName)
	require.Nil(t, isolated)
	isolated, _ = PreferredPoolChannel(ctx(100, "session-a"), "another-model")
	require.Nil(t, isolated)
	ClearPoolAffinity(next)
	bound, _ = PreferredPoolChannel(ctx(100, "session-a"), modelName)
	require.Nil(t, bound)
	AppendUsedChannel(c, second.Id)
	last, _, err2 := CacheGetRandomSatisfiedChannel(param)
	require.NoError(t, err2)
	require.Nil(t, last)
}