package service

import (
	"context"
	"errors"
	"fmt"
	"github.com/QuantumNous/new-api/common"
	"github.com/QuantumNous/new-api/constant"
	"github.com/QuantumNous/new-api/model"
	perfmetrics "github.com/QuantumNous/new-api/pkg/perf_metrics"
	apiTypes "github.com/QuantumNous/new-api/relaykit/types"
	"github.com/QuantumNous/new-api/setting/operation_setting"
	"github.com/gin-gonic/gin"
	"github.com/stretchr/testify/require"
	"math"
	"net/http/httptest"
	"slices"
	"testing"
	"time"
)

func TestTokenRoutingSingleChannelRetryCooldownAndRecovery(t *testing.T) {
	db := setupChannelSelectAutoGroupsTest(t)
	createChannelSelectAutoGroupsChannel(t, db, 6911, "default", "retry-model")
	createChannelSelectAutoGroupsChannel(t, db, 6912, "vip", "retry-model")
	model.InitChannelCache()
	ResetGroupAvailability()
	t.Cleanup(func() { ResetGroupAvailability(); _ = poolCooldowns().Purge() })
	makeCtx := func(token int) *gin.Context {
		c := setupGroupPoolContext(t)
		c.Request = httptest.NewRequest("POST", "/v1/responses", nil)
		c.Set("id", 99)
		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}}})
		return c
	}
	selectID := func(c *gin.Context) int {
		t.Helper()
		ch, _, err := selectTokenRoutingPool(&RetryParam{Ctx: c, ModelName: "retry-model", Retry: common.GetPointer(0)})
		require.NoError(t, err)
		require.NotNil(t, ch)
		AppendUsedChannel(c, ch.Id)
		return ch.Id
	}
	c := makeCtx(991)
	require.Equal(t, 6911, selectID(c))
	RecordPoolRetryableFailure(c, "retry-model", "default")
	require.Equal(t, 6911, selectID(c), "a one-channel group must actually get its second attempt")
	RecordPoolRetryableFailure(c, "retry-model", "default")
	require.Equal(t, 6912, selectID(c))
	require.Equal(t, 6912, selectID(makeCtx(991)))
	require.Equal(t, 6911, selectID(makeCtx(992)))
	key := poolCooldownKey(c, "retry-model", "default")
	state, found, err := poolCooldowns().Get(key)
	require.NoError(t, err)
	require.True(t, found)
	require.InDelta(t, int64((10*time.Minute)/time.Millisecond), state.Until-time.Now().UnixMilli(), 2000)
	require.NoError(t, poolCooldowns().SetWithTTL(key, poolCooldown{Until: time.Now().Add(-time.Second).UnixMilli()}, 0))
	trial := makeCtx(991)
	require.Equal(t, 6911, selectID(trial))
	require.Equal(t, 6912, selectID(makeCtx(991)), "concurrent request waits for recovery probe")
	RecordPoolRetryableFailure(trial, "retry-model", "default")
	require.Equal(t, 6912, selectID(trial), "failed probe switches without another primary attempt")
	require.NoError(t, poolCooldowns().SetWithTTL(key, poolCooldown{Until: time.Now().Add(-time.Second).UnixMilli()}, 0))
	recovery := makeCtx(991)
	require.Equal(t, 6911, selectID(recovery))
	RecordPoolSuccess(recovery, "retry-model", "default")
	_, found, err = poolCooldowns().Get(key)
	require.NoError(t, err)
	require.False(t, found)
	next := makeCtx(991)
	require.Equal(t, 6911, selectID(next))
	RecordPoolRetryableFailure(next, "retry-model", "default")
	require.Equal(t, 6911, selectID(next), "recovered group returns to two-attempt behavior")
	mixed := makeCtx(993)
	require.Equal(t, 6911, selectID(mixed))
	RecordPoolRetryableFailure(mixed, "retry-model", "default")
	require.Equal(t, 6911, selectID(mixed))
	require.Equal(t, 6912, selectID(mixed), "the one same-group retry must not repeat when a different error triggers another selection")
}

func TestAutomaticGroupPreferenceRanking(t *testing.T) {
	original := []automaticCandidate{{group: "cheap", cost: .1, ttft: 900, success: 90, ratio: .2}, {group: "fast", cost: .2, ttft: 100, success: 95, ratio: .3}, {group: "stable", cost: .3, ttft: 400, success: 99, ratio: .4}, {group: "unknown", cost: math.Inf(1), ttft: math.Inf(1), success: -1, ratio: .01}}
	for preference, want := range map[string]string{"cost": "unknown", "speed": "fast", "stable": "stable"} {
		t.Run(preference, func(t *testing.T) {
			items := slices.Clone(original)
			sortAutomaticCandidates(items, preference)
			require.Equal(t, want, items[0].group)
			if preference != "cost" {
				require.Equal(t, "unknown", items[len(items)-1].group)
			}
		})
	}
}

func TestAutomaticGroupUsesAllAuthorizedModelGroups(t *testing.T) {
	db := setupChannelSelectAutoGroupsTest(t)
	for _, fixture := range []struct {
		id          int
		group, name string
	}{{6921, "default", "smart-model"}, {6922, "vip", "smart-model"}, {6923, "secret", "smart-model"}, {6924, "default", "another-model"}} {
		createChannelSelectAutoGroupsChannel(t, db, fixture.id, fixture.group, fixture.name)
	}
	model.InitChannelCache()
	groupModelCostSnapshot.Lock()
	oldCosts, oldTime := groupModelCostSnapshot.values, groupModelCostSnapshot.updated
	cheap, expensive := 0.1, 0.3
	groupModelCostSnapshot.values = map[string]map[string]model.GroupModelCost{"default": {"smart-model": {USDPerMillion: &expensive}}, "vip": {"smart-model": {USDPerMillion: &cheap}}}
	groupModelCostSnapshot.updated = time.Now()
	groupModelCostSnapshot.Unlock()
	automaticMetricCache.Set("smart-model", perfmetrics.QueryResult{})
	t.Cleanup(func() {
		groupModelCostSnapshot.Lock()
		groupModelCostSnapshot.values, groupModelCostSnapshot.updated = oldCosts, oldTime
		groupModelCostSnapshot.Unlock()
		automaticMetricCache.Purge()
	})
	makeCtx := func() *gin.Context {
		c := setupGroupPoolContext(t)
		c.Request = httptest.NewRequest("POST", "/v1/responses", nil)
		c.Set(model.TokenRoutingContextKey, &model.TokenRoutingPolicy{Version: 1, Mode: "smart", AutoPreference: "cost"})
		return c
	}
	result := ResolveRequestPool(makeCtx(), "smart-model")
	require.Equal(t, []model.TokenGroupEntry{{Group: "default", Priority: 0}, {Group: "vip", Priority: 1}}, result)
	list := ResolveRequestPool(makeCtx(), "")
	require.Len(t, list, 2)
	for _, entry := range list {
		require.NotEqual(t, "secret", entry.Group)
	}
	require.Empty(t, ResolveRequestPool(makeCtx(), "unknown-model"))
}

func TestPoolFailoverRespectsRetrySafetyBoundaries(t *testing.T) {
	for _, code := range []int{400, 401, 429, 500, 502, 503, 504, 524} {
		t.Run(fmt.Sprint(code), func(t *testing.T) {
			c := setupGroupPoolContext(t)
			c.Request = httptest.NewRequest("POST", "/v1/responses", nil)
			c.Set(model.TokenRoutingContextKey, &model.TokenRoutingPolicy{Version: 1, Mode: "smart"})
			err := apiTypes.NewErrorWithStatusCode(errors.New("test upstream"), apiTypes.ErrorCodeBadResponseStatusCode, code)
			require.Equal(t, code == 500 || code == 502 || code == 503, IsPoolFailoverError(c, err))
			ctx, cancel := context.WithCancel(c.Request.Context())
			cancel()
			c.Request = c.Request.WithContext(ctx)
			require.False(t, IsPoolFailoverError(c, err))
		})
	}
	c := setupGroupPoolContext(t)
	c.Request = httptest.NewRequest("POST", "/v1/responses", nil)
	c.Set(model.TokenRoutingContextKey, &model.TokenRoutingPolicy{Version: 1, Mode: "smart"})
	_, err := c.Writer.WriteString("already streaming")
	require.NoError(t, err)
	failure := apiTypes.NewErrorWithStatusCode(errors.New("test upstream"), apiTypes.ErrorCodeBadResponseStatusCode, 500)
	require.False(t, IsPoolFailoverError(c, failure))
	require.Equal(t, "stop", DecideRelayRetry(c, failure, 5).Action)
}

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)
}