package service

import (
	"context"
	"github.com/QuantumNous/new-api/common"
	"github.com/QuantumNous/new-api/model"
	"github.com/gin-gonic/gin"
	"github.com/stretchr/testify/require"
	"net/http/httptest"
	"testing"
	"time"
)

func TestPoolWithoutAvailableBackupDoesNotCoolDown(t *testing.T) {
	db := setupChannelSelectAutoGroupsTest(t)
	createChannelSelectAutoGroupsChannel(t, db, 7801, "default", "m")
	createChannelSelectAutoGroupsChannel(t, db, 7802, "vip", "other")
	model.InitChannelCache()
	ResetGroupAvailability()
	t.Cleanup(func() { ResetGroupAvailability(); _ = poolCooldowns().Purge() })
	makeCtx := func() *gin.Context {
		c := setupGroupPoolContext(t)
		c.Request = httptest.NewRequest("POST", "/v1/responses", nil)
		c.Set("token_id", 7801)
		c.Set(model.TokenRoutingContextKey, &model.TokenRoutingPolicy{Version: 1, Mode: "custom", Groups: []model.TokenGroupEntry{{Group: "default"}, {Group: "vip", Priority: 1}}})
		return c
	}
	c := makeCtx()
	RecordPoolRetryableFailure(c, "m", "default")
	RecordPoolRetryableFailure(c, "m", "default")
	_, found, err := poolCooldowns().Get(poolCooldownKey(c, "m", "default"))
	require.NoError(t, err)
	require.False(t, found)
	// An old cooldown and global health state must not strand a single eligible group.
	require.NoError(t, poolCooldowns().SetWithTTL(poolCooldownKey(c, "m", "default"), poolCooldown{Until: time.Now().Add(10 * time.Minute).UnixMilli()}, 0))
	for i := 0; i < 4; i++ {
		RecordGroupOutcome("default", false)
	}
	next := makeCtx()
	ch, group, err := selectTokenRoutingPool(&RetryParam{Ctx: next, ModelName: "m", Retry: common.GetPointer(0)})
	require.NoError(t, err)
	require.NotNil(t, ch)
	require.Equal(t, "default", group)
	AppendUsedChannel(next, ch.Id)
	RecordPoolRetryableFailure(next, "m", "default")
	ch, _, err = selectTokenRoutingPool(&RetryParam{Ctx: next, ModelName: "m", Retry: common.GetPointer(1)})
	require.NoError(t, err)
	require.NotNil(t, ch)
	AppendUsedChannel(next, ch.Id)
	RecordPoolRetryableFailure(next, "m", "default")
	ch, _, err = selectTokenRoutingPool(&RetryParam{Ctx: next, ModelName: "m", Retry: common.GetPointer(2)})
	require.NoError(t, err)
	require.Nil(t, ch, "do not loop indefinitely inside one request")
}
func TestAutomaticSpeedUsesLatestExactModelStreamingTest(t *testing.T) {
	db := setupChannelSelectAutoGroupsTest(t)
	require.NoError(t, db.AutoMigrate(&model.Log{}))
	previous := model.LOG_DB
	model.LOG_DB = db
	t.Cleanup(func() { model.LOG_DB = previous })
	now := time.Now().Unix()
	samples := []model.Log{
		{ChannelId: 1, Type: 2, TokenId: 0, Content: "模型测试", ModelName: "m", IsStream: true, CreatedAt: now - 20, Other: `{"frt":500}`},
		{ChannelId: 1, Type: 2, TokenId: 0, Content: "模型测试", ModelName: "m", IsStream: true, CreatedAt: now - 10, Other: `{"frt":60000}`},
		{ChannelId: 1, Type: 2, TokenId: 9, Content: "模型测试", ModelName: "m", IsStream: true, CreatedAt: now, Other: `{"frt":1}`},
		{ChannelId: 2, Type: 2, TokenId: 0, Content: "模型测试", ModelName: "m", IsStream: true, CreatedAt: now - 5, Other: `{"frt":1200}`},
		{ChannelId: 2, Type: 2, TokenId: 0, Content: "模型测试", ModelName: "other", IsStream: true, CreatedAt: now, Other: `{"frt":10}`},
		{ChannelId: 3, Type: 2, TokenId: 0, Content: "模型测试", ModelName: "m", IsStream: false, CreatedAt: now, Other: `{"frt":10}`},
		{ChannelId: 4, Type: 2, TokenId: 0, Content: "模型测试", ModelName: "m", IsStream: true, CreatedAt: now - 2*3600, Other: `{"frt":10}`},
	}
	require.NoError(t, db.Create(&samples).Error)
	speeds := latestModelTestSpeeds(context.Background(), "m")
	require.Equal(t, map[int]float64{1: 60000, 2: 1200}, speeds)
}

func TestProbeAndStreamingCooldownsAreIsolated(t *testing.T) {
	db := setupChannelSelectAutoGroupsTest(t)
	createChannelSelectAutoGroupsChannel(t, db, 7901, "default", "m")
	createChannelSelectAutoGroupsChannel(t, db, 7902, "vip", "m")
	model.InitChannelCache()
	t.Cleanup(func() { _ = poolCooldowns().Purge() })
	makeCtx := func(stream bool) *gin.Context {
		c := setupGroupPoolContext(t)
		c.Request = httptest.NewRequest("POST", "/v1/responses", nil)
		c.Set("token_id", 7901)
		c.Set("pool_request_stream", stream)
		c.Set(model.TokenRoutingContextKey, &model.TokenRoutingPolicy{Version: 1, Mode: "custom", Groups: []model.TokenGroupEntry{{Group: "default"}, {Group: "vip", Priority: 1}}})
		return c
	}
	probe := makeCtx(false)
	RecordPoolRetryableFailure(probe, "m", "default")
	RecordPoolRetryableFailure(probe, "m", "default")
	require.True(t, poolGroupSkipped(makeCtx(false), "m", "default"))
	require.False(t, poolGroupSkipped(makeCtx(true), "m", "default"))
}