package service

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

	"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"
	"github.com/QuantumNous/new-api/setting"
	"github.com/QuantumNous/new-api/setting/operation_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 TestPoolAffinityManualRoutingSwitch(t *testing.T) {
	for _, change := range []string{"priority", "model_rule", "blacklist"} {
		t.Run(change, func(t *testing.T) {
			db := setupChannelSelectAutoGroupsTest(t)
			const modelName = "manual-switch-model"
			createChannelSelectAutoGroupsChannel(t, db, 6201, "default", modelName)
			createChannelSelectAutoGroupsChannel(t, db, 6202, "vip", modelName)
			model.InitChannelCache()
			ResetGroupAvailability()
			t.Cleanup(ResetGroupAvailability)
			previousRedis := common.RedisEnabled
			common.RedisEnabled = false
			t.Cleanup(func() { common.RedisEnabled = previousRedis })
			require.NoError(t, poolBindings().Purge())
			t.Cleanup(func() { _ = poolBindings().Purge() })
			config := operation_setting.GetChannelAffinitySetting()
			previousConfig := *config
			*config = operation_setting.ChannelAffinitySetting{Enabled: true, DefaultTTLSeconds: 3600, SessionMode: "strict", Rules: []operation_setting.ChannelAffinityRule{{Name: "switch-test", ModelRegex: []string{".*"}, KeySources: []operation_setting.ChannelAffinityKeySource{{Type: "request_header", Key: "Session-Id"}}, SessionMode: "strict"}}}
			t.Cleanup(func() { *config = previousConfig })
			previousLogDB := model.LOG_DB
			model.LOG_DB = db
			t.Cleanup(func() { model.LOG_DB = previousLogDB })
			require.NoError(t, db.AutoMigrate(&model.Log{}))
			ctx := func(switched bool) *gin.Context {
				c := setupGroupPoolContext(t)
				c.Request = httptest.NewRequest("POST", "/v1/responses", nil)
				c.Request.Header.Set("Session-Id", "same-session")
				c.Set("id", 100)
				c.Set("token_id", 100)
				c.Set("original_model", modelName)
				policy := &model.TokenRoutingPolicy{Version: 1, Mode: "custom", Groups: []model.TokenGroupEntry{{Group: "default", Priority: 0}, {Group: "vip", Priority: 1}}}
				if change == "blacklist" {
					policy.Mode = "smart"
					policy.Groups = nil
					automaticMetricCache.Set(modelName, perfmetrics.QueryResult{})
				}
				if switched {
					if change == "priority" {
						policy.Groups[0].Priority = 1
						policy.Groups[1].Priority = 0
					} else if change == "blacklist" {
						policy.ExcludedGroups = []string{"default"}
					} else {
						policy.ModelRules = []model.ModelRoutingRule{{Model: modelName, Groups: []model.TokenGroupEntry{{Group: "vip", Priority: 0}}}}
					}
				}
				c.Set(model.TokenRoutingContextKey, policy)
				return c
			}
			old := ctx(false)
			ch, _, selectErr := SelectChannelForRequest(old, modelName, &RetryParam{Ctx: old, TokenGroup: "default", ModelName: modelName, Retry: common.GetPointer(0)})
			require.Nil(t, selectErr)
			require.NotNil(t, ch)
			assert.Equal(t, 6201, ch.Id)
			old.Set("channel_id", ch.Id)
			common.SetContextKey(old, constant.ContextKeyAutoGroup, "default")
			RequestPolicy(old).Successful = true
			RecordPoolAffinity(old)
			unchanged, _ := PreferredPoolChannel(ctx(false), modelName)
			require.NotNil(t, unchanged)
			assert.Equal(t, 6201, unchanged.Id)

			details := model.NewLogOther()
			appendRequestDetails(old, nil, details)
			body, err := common.Marshal(details.Snapshot())
			require.NoError(t, err)
			log := model.Log{UserId: 100, TokenId: 100, ModelName: modelName, CorrelationKey: old.GetString("log_correlation_key"), Type: model.LogTypeConsume, CreatedAt: time.Now().Unix(), ChannelId: 6201, Group: "default", Other: string(body)}
			require.NoError(t, db.Create(&log).Error)

			// The edited policy must take precedence over the existing binding.
			switched := ctx(true)
			selected, group, selectErr := SelectChannelForRequest(switched, modelName, &RetryParam{Ctx: switched, TokenGroup: "default", ModelName: modelName, Retry: common.GetPointer(0)})
			require.Nil(t, selectErr)
			require.NotNil(t, selected)
			assert.Equal(t, 6202, selected.Id)
			assert.Equal(t, "vip", group)
			switched.Set("channel_id", selected.Id)
			common.SetContextKey(switched, constant.ContextKeyAutoGroup, group)
			RequestPolicy(switched).Successful = true
			RecordPoolAffinity(switched)
			unchanged, _ = PreferredPoolChannel(ctx(true), modelName)
			require.NotNil(t, unchanged)
			assert.Equal(t, 6202, unchanged.Id)

			// A request started before the edit may finish after the new one.
			RecordPoolAffinity(old)
			stale, _ := PreferredPoolChannel(ctx(true), modelName)
			assert.Nil(t, stale)

			// Cache expiry/restart must not resurrect the old successful log.
			require.NoError(t, poolBindings().Purge())
			stale, _ = PreferredPoolChannel(ctx(true), modelName)
			assert.Nil(t, stale)

			// A matching historical policy remains recoverable after cache expiry.
			require.NoError(t, poolBindings().Purge())
			recovered, _ := PreferredPoolChannel(ctx(false), modelName)
			require.NotNil(t, recovered)
			assert.Equal(t, 6201, recovered.Id)

			// Legacy bindings carry no policy and must be selected afresh once.
			require.NoError(t, poolBindings().SetWithTTL(old.GetString("pool_affinity_key"), poolBinding{Channel: 6201, Group: "default"}, time.Hour))
			legacy, _ := PreferredPoolChannel(ctx(false), modelName)
			assert.Nil(t, legacy)
		})
	}
}

func TestAutomaticPoolKeepsHealthySession(t *testing.T) {
	for _, scenario := range []string{"cheaper_recovered", "ranking_changed", "candidate_added", "degraded", "cooldown", "permission_removed", "channel_disabled", "policy_changed", "new_session", "session_off", "custom_recovery"} {
		t.Run(scenario, func(t *testing.T) {
			db := setupChannelSelectAutoGroupsTest(t)
			const modelName = "smart-sticky-model"
			createChannelSelectAutoGroupsChannel(t, db, 6301, "default", modelName)
			createChannelSelectAutoGroupsChannel(t, db, 6302, "vip", modelName)
			model.InitChannelCache()
			oldRedis, oldLogDB := common.RedisEnabled, model.LOG_DB
			common.RedisEnabled, model.LOG_DB = false, nil
			config := operation_setting.GetChannelAffinitySetting()
			oldConfig := *config
			*config = operation_setting.ChannelAffinitySetting{Enabled: true, DefaultTTLSeconds: 3600, SessionMode: "prefer", Rules: []operation_setting.ChannelAffinityRule{{Name: "smart-sticky", ModelRegex: []string{".*"}, KeySources: []operation_setting.ChannelAffinityKeySource{{Type: "request_header", Key: "Session-Id"}}, SessionMode: "prefer"}}}
			require.NoError(t, poolBindings().Purge())
			require.NoError(t, poolCooldowns().Purge())
			automaticMetricCache.Set(modelName, perfmetrics.QueryResult{})
			t.Cleanup(func() {
				*config = oldConfig
				common.RedisEnabled, model.LOG_DB = oldRedis, oldLogDB
				_ = poolBindings().Purge()
				_ = poolCooldowns().Purge()
				automaticMetricCache.Purge()
			})
			ctx := func() *gin.Context {
				c := setupGroupPoolContext(t)
				c.Request = httptest.NewRequest("POST", "/v1/responses", nil)
				c.Request.Header.Set("Session-Id", "same-session")
				c.Set("id", 101)
				c.Set("token_id", 102)
				c.Set("original_model", modelName)
				policy := &model.TokenRoutingPolicy{Version: 1, Mode: "smart", AutoPreference: "cost", SkipDegradedNonAstra: true}
				if scenario == "custom_recovery" {
					policy.Mode = "custom"
					policy.Groups = []model.TokenGroupEntry{{Group: "default", Priority: 0}, {Group: "vip", Priority: 1}}
				}
				c.Set(model.TokenRoutingContextKey, policy)
				c.Set("pool_group_fingerprints", map[string]AstraSamplingHistory{})
				c.Set("pool_group_visions", map[string]DetectionHistory{})
				return c
			}
			// A temporary failure sends the session to the more expensive group.
			first := ctx()
			cheapKey := poolCooldownKey(first, modelName, "default")
			require.NoError(t, poolCooldowns().SetWithTTL(cheapKey, poolCooldown{Until: time.Now().Add(time.Minute).UnixMilli()}, 0))
			if scenario == "candidate_added" {
				require.NoError(t, setting.UpdateUserUsableGroupsByJSONString(`{"vip":"VIP"}`))
			}
			ch, group, selectErr := SelectChannelForRequest(first, modelName, &RetryParam{Ctx: first, ModelName: modelName, Retry: common.GetPointer(0)})
			require.Nil(t, selectErr)
			require.NotNil(t, ch)
			require.Equal(t, "vip", group)
			first.Set("channel_id", ch.Id)
			RequestPolicy(first).Successful = true
			RecordPoolAffinity(first)
			require.NoError(t, poolCooldowns().SetWithTTL(cheapKey, poolCooldown{Until: time.Now().Add(-time.Second).UnixMilli()}, 0))
			next := ctx()
			want := "vip"
			switch scenario {
			case "ranking_changed":
				// The initial pool already placed default first. Reverse it to prove ranking changes are irrelevant.
				require.NoError(t, ratio_setting.UpdateGroupRatioByJSONString(`{"default":4,"vip":3}`))
			case "candidate_added":
				require.NoError(t, setting.UpdateUserUsableGroupsByJSONString(`{"default":"Default","vip":"VIP"}`))
			case "degraded":
				next.Set("pool_group_visions", map[string]DetectionHistory{"vip": {Quality: "degraded"}})
				want = "default"
			case "cooldown":
				require.NoError(t, poolCooldowns().SetWithTTL(poolCooldownKey(next, modelName, "vip"), poolCooldown{Until: time.Now().Add(time.Minute).UnixMilli()}, 0))
				want = "default"
			case "permission_removed":
				require.NoError(t, setting.UpdateUserUsableGroupsByJSONString(`{"default":"Default"}`))
				want = "default"
			case "channel_disabled":
				require.NoError(t, db.Model(&model.Channel{}).Where("id = ?", 6302).Update("status", common.ChannelStatusAutoDisabled).Error)
				model.InitChannelCache()
				want = "default"
			case "policy_changed":
				RequestRoutingPolicy(next).AutoPreference = "stable"
				want = "default"
			case "new_session":
				next.Request.Header.Set("Session-Id", "new-session")
				want = "default"
			case "session_off":
				config.Rules[0].SessionMode = "off"
				config.SessionMode = "off"
				want = "default"
			case "custom_recovery":
				want = "default"
			}
			selected, selectedGroup, selectErr := SelectChannelForRequest(next, modelName, &RetryParam{Ctx: next, ModelName: modelName, Retry: common.GetPointer(0)})
			require.Nil(t, selectErr)
			require.NotNil(t, selected)
			assert.Equal(t, want, selectedGroup)
			if want == "vip" {
				assert.Equal(t, 6302, selected.Id)
				state, found, err := poolCooldowns().Get(cheapKey)
				require.NoError(t, err)
				require.True(t, found)
				assert.Empty(t, state.ProbeID, "keeping a healthy route must not claim a cheaper group's recovery probe")
			}
		})
	}
}