package service

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

	"github.com/QuantumNous/new-api/common"
	"github.com/QuantumNous/new-api/model"
	fingerprint "github.com/QuantumNous/new-api/pkg/model_fingerprint"
	"github.com/gin-gonic/gin"
	"github.com/stretchr/testify/require"
)

func TestQualitySwitchesAndFailedReadCache(t *testing.T) {
	db := setupChannelSelectAutoGroupsTest(t)
	require.NoError(t, db.AutoMigrate(&model.Option{}, &model.SystemTask{}, &model.DetectionRun{}))
	createChannelSelectAutoGroupsChannel(t, db, 6501, "default", fingerprint.Model)
	old := time.Now().Add(-8 * time.Hour).Unix()
	result, err := common.Marshal(fingerprint.Result{Engine: fingerprint.Version, Status: "mismatch", ValidSamples: 3, Closest: "gpt-6-sol"})
	require.NoError(t, err)
	run := model.DetectionRun{Result: string(result), ID: "old-valid", RequestKey: "old-valid", ChannelID: 6501, Kind: "fingerprint", Status: "succeeded", Settled: true, Quality: "degraded", CreatedAt: old}
	require.NoError(t, db.Create(&run).Error)
	require.NoError(t, db.Create(&model.DetectionRun{ID: "failed", RequestKey: "failed", ChannelID: 6501, Kind: "pelican", Status: "failed", Settled: true, CreatedAt: time.Now().Unix()}).Error)
	groupDegradationSnapshot.latest = nil
	InvalidateGroupQualityCache()
	t.Cleanup(func() {
		groupDegradationSnapshot.latest = nil
		groupDegradationSnapshot.vision = nil
		groupDegradationSnapshot.ciallo = nil
		InvalidateGroupQualityCache()
	})
	ctx := func(allow, child bool) *gin.Context {
		c, _ := gin.CreateTestContext(httptest.NewRecorder())
		c.Request = httptest.NewRequest("POST", "/v1/responses", nil)
		c.Set(model.TokenRoutingContextKey, &model.TokenRoutingPolicy{Version: 1, Mode: "smart", AllowDegraded: allow, SkipDegradedNonAstra: child})
		return c
	}
	for _, allow := range []bool{false, true} {
		for _, child := range []bool{false, true} {
			require.Equal(t, !allow, requestGroupDegraded(ctx(allow, child), fingerprint.Model, "default"))
			require.Equal(t, !allow && child, requestGroupDegraded(ctx(allow, child), "gpt-6-sol", "default"))
		}
	}
	recent, err := RecentDegradedGroups(context.Background())
	require.NoError(t, err)
	require.False(t, recent["default"])
	require.NoError(t, db.Migrator().DropTable(&model.DetectionRun{}))
	InvalidateGroupQualityCache()
	require.True(t, requestGroupDegraded(ctx(false, false), fingerprint.Model, "default"), "read failure retains valid conclusion")
	groupDegradationSnapshot.latest = nil
	InvalidateGroupQualityCache()
	cold := ctx(false, false)
	require.True(t, requestGroupDegraded(cold, fingerprint.Model, "default"))
	require.True(t, cold.GetBool("pool_degradation_unavailable"))
	require.NoError(t, db.AutoMigrate(&model.DetectionRun{}))
	InvalidateGroupQualityCache()
	require.False(t, requestGroupDegraded(ctx(false, false), fingerprint.Model, "default"), "successful empty history remains eligible")
}

func TestAbandonedPoolProbeIsReleasedWithoutClearingFailures(t *testing.T) {
	db := setupChannelSelectAutoGroupsTest(t)
	createChannelSelectAutoGroupsChannel(t, db, 7001, "vip", "m")
	model.InitChannelCache()
	c := setupGroupPoolContext(t)
	c.Request = httptest.NewRequest("POST", "/v1/responses", nil)
	c.Set("id", 701)
	c.Set("token_id", 701)
	key := poolCooldownKey(c, "m", "default")
	require.NoError(t, poolCooldowns().SetWithTTL(key, poolCooldown{Until: time.Now().Add(-time.Second).UnixMilli()}, 0))
	t.Cleanup(func() { _, _ = poolCooldowns().DeleteMany([]string{key}) })
	require.True(t, claimPoolGroup(c, "m", "default"))
	require.True(t, poolGroupSkipped(c, "m", "default"))
	ReleasePoolProbes(c)
	require.False(t, poolGroupSkipped(c, "m", "default"))
	require.True(t, claimPoolGroup(c, "m", "default"))
	c.Set(model.TokenRoutingContextKey, &model.TokenRoutingPolicy{Version: 1, Mode: "smart"})
	RecordPoolRetryableFailure(c, "m", "default")
	ReleasePoolProbes(c)
	value, found, err := poolCooldowns().Get(key)
	require.NoError(t, err)
	require.True(t, found)
	require.Greater(t, value.Until, time.Now().UnixMilli())
}

func TestManualPoolNonAstraDegradationSwitch(t *testing.T) {
	for _, mode := range []string{"inherit", "custom"} {
		for _, enabled := range []bool{false, true} {
			for _, nonAstra := range []bool{false, true} {
				c, _ := gin.CreateTestContext(httptest.NewRecorder())
				c.Set(model.TokenRoutingContextKey, &model.TokenRoutingPolicy{Version: 1, Mode: mode, SkipDegradedGroups: enabled, SkipDegradedNonAstra: nonAstra})
				c.Set("pool_group_fingerprints", map[string]AstraSamplingHistory{"default": {Status: "mismatch", ValidSamples: 3, Closest: "gpt-6-sol"}})
				c.Set("pool_group_pelicans", map[string]DetectionHistory{"default": {Quality: "degraded"}})
				require.Equal(t, enabled, requestGroupDegraded(c, fingerprint.Model, "default"))
				require.Equal(t, enabled && nonAstra, requestGroupDegraded(c, "gpt-6-sol", "default"))
				require.False(t, requestGroupDegraded(c, "gpt-6-sol", "untested"))
			}
		}
	}
}