package service

import (
	"fmt"
	"os"
	"strings"
	"testing"
	"time"

	"github.com/QuantumNous/new-api/common"
	"github.com/QuantumNous/new-api/constant"
	"github.com/QuantumNous/new-api/model"
	fingerprint "github.com/QuantumNous/new-api/pkg/model_fingerprint"
	"github.com/QuantumNous/new-api/setting/operation_setting"
	"github.com/gin-gonic/gin"
	"github.com/glebarez/sqlite"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
	"gorm.io/driver/mysql"
	"gorm.io/driver/postgres"
	"gorm.io/gorm"
	"gorm.io/gorm/schema"
	"net/http/httptest"
)

func TestGroupDegradationDatabaseMatrix(t *testing.T) {
	for _, dialect := range []string{"sqlite", "mysql", "postgres"} {
		t.Run(dialect, func(t *testing.T) {
			var driver gorm.Dialector
			switch dialect {
			case "sqlite":
				driver = sqlite.Open(":memory:")
			case "mysql":
				if os.Getenv("TEST_MYSQL_DSN") == "" {
					t.Skip("TEST_MYSQL_DSN unset")
				}
				driver = mysql.Open(os.Getenv("TEST_MYSQL_DSN"))
			case "postgres":
				if os.Getenv("TEST_POSTGRES_DSN") == "" {
					t.Skip("TEST_POSTGRES_DSN unset")
				}
				driver = postgres.Open(os.Getenv("TEST_POSTGRES_DSN"))
			}
			db, err := gorm.Open(driver, &gorm.Config{NamingStrategy: schema.NamingStrategy{TablePrefix: fmt.Sprintf("degradation_%d_", time.Now().UnixNano())}})
			require.NoError(t, err)
			oldDB := model.DB
			oldType, oldLogType := common.MainDatabaseType(), common.LogDatabaseType()
			model.DB = db
			common.SetDatabaseTypes(common.DatabaseType(dialect), oldLogType)
			rows := []any{&model.Option{}, &model.Channel{}, &model.SystemTask{}, &model.DetectionSchedule{}, &model.DetectionRun{}, &model.DetectionCall{}, &model.Token{}}
			require.NoError(t, db.AutoMigrate(rows...))
			t.Cleanup(func() {
				for _, row := range rows {
					require.NoError(t, db.Migrator().DropTable(row))
				}
				model.DB = oldDB
				common.SetDatabaseTypes(oldType, oldLogType)
				groupDegradationSnapshot.updated = time.Time{}
				groupDegradationSnapshot.latest = nil
				groupDegradationSnapshot.vision = nil
				groupDegradationSnapshot.ciallo = nil
				conn, err := db.DB()
				require.NoError(t, err)
				require.NoError(t, conn.Close())
			})
			var version string
			query := "select version()"
			if dialect == "sqlite" {
				query = "select sqlite_version()"
			}
			require.NoError(t, db.Raw(query).Scan(&version).Error)
			t.Log("database version:", version)
			now := time.Now().Unix()
			require.NoError(t, db.Create(&model.Channel{Id: 1, Group: "default", Models: fingerprint.Model, Status: 1, Key: "test-only"}).Error)
			body, err := common.Marshal(AstraSamplingState{Rounds: []AstraSamplingRound{{ID: "mismatch", ChannelID: 1, StartedAt: now - 30, FinishedAt: now - 10, Result: fingerprint.Result{Engine: fingerprint.Version, Status: "mismatch", Closest: "gpt-5.5", ValidSamples: 2}}}})
			require.NoError(t, err)
			require.NoError(t, db.Create(&model.SystemTask{TaskID: "degradation-test", Type: AstraSamplingTaskType, Status: model.SystemTaskStatusSucceeded, CreatedAt: now - 30, State: string(body)}).Error)
			policy := &model.TokenRoutingPolicy{Version: 1, Mode: "custom", SkipDegradedGroups: true, Groups: []model.TokenGroupEntry{{Group: "default", Priority: 0}}}
			encoded, err := common.Marshal(policy)
			require.NoError(t, err)
			token := model.Token{UserId: 1, Key: "degradation-test-key", RoutingPolicy: string(encoded)}
			require.NoError(t, db.Create(&token).Error)
			require.NoError(t, db.First(&token, token.Id).Error)
			stored, err := token.GetRoutingPolicy()
			require.NoError(t, err)
			require.True(t, stored.SkipDegradedGroups)
			ctx := func() *gin.Context {
				c, _ := gin.CreateTestContext(httptest.NewRecorder())
				c.Request = httptest.NewRequest("POST", "/v1/responses", nil)
				c.Set(model.TokenRoutingContextKey, stored)
				return c
			}
			groupDegradationSnapshot.updated = time.Time{}
			assert.True(t, requestGroupDegraded(ctx(), fingerprint.Model, "default"))
			smart := ctx()
			smartPolicy := &model.TokenRoutingPolicy{Version: 1, Mode: "smart"}
			smart.Set(model.TokenRoutingContextKey, smartPolicy)
			assert.True(t, requestGroupDegraded(smart, fingerprint.Model, "default"))
			smartPolicy.AllowDegraded = true
			assert.False(t, requestGroupDegraded(smart, fingerprint.Model, "default"))
			smartPolicy.AllowDegraded = false
			assert.False(t, requestGroupDegraded(smart, "unmeasured-model", "default"))
			manualResult, err := common.Marshal(fingerprint.Result{Engine: fingerprint.Version, Status: "match", Closest: fingerprint.Model, ValidSamples: 2})
			require.NoError(t, err)
			manual := model.DetectionRun{ID: "manual-recovery", RequestKey: "manual-recovery", ChannelID: 1, Kind: "fingerprint", UserID: 1, CreatedAt: now, Settled: true, Status: "succeeded", Result: string(manualResult)}
			require.NoError(t, db.Create(&manual).Error)
			groupDegradationSnapshot.updated = time.Time{}
			assert.False(t, requestGroupDegraded(ctx(), fingerprint.Model, "default"))
			pelican := model.DetectionRun{ID: "pelican-degraded", RequestKey: "pelican-degraded", ChannelID: 1, Kind: "pelican", CreatedAt: now, Settled: true, Status: "succeeded", Quality: "degraded"}
			require.NoError(t, db.Create(&pelican).Error)
			groupDegradationSnapshot.updated = time.Time{}
			assert.False(t, requestGroupDegraded(ctx(), fingerprint.Model, "default"))
			require.NoError(t, db.Model(&pelican).Update("quality", "normal").Error)
			groupDegradationSnapshot.updated = time.Time{}
			assert.False(t, requestGroupDegraded(ctx(), fingerprint.Model, "default"))

			vision := model.DetectionRun{ID: "vision-degraded", RequestKey: "vision-degraded", ChannelID: 1, Kind: "vision", CreatedAt: now - 5, Settled: true, Status: "succeeded", Quality: "degraded", Result: `{"closest":"石头"}`}
			require.NoError(t, db.Create(&vision).Error)
			InvalidateGroupQualityCache()
			assert.True(t, requestGroupDegraded(ctx(), fingerprint.Model, "default"), "fingerprint recovery must not clear a failed vision test")
			recent, err := RecentDegradedGroups(ctx().Request.Context())
			require.NoError(t, err)
			assert.True(t, recent["default"])
			failedVision := model.DetectionRun{ID: "vision-error", RequestKey: "vision-error", ChannelID: 1, Kind: "vision", CreatedAt: now - 1, Settled: true, Status: "failed", Quality: "normal"}
			require.NoError(t, db.Create(&failedVision).Error)
			InvalidateGroupQualityCache()
			assert.True(t, requestGroupDegraded(ctx(), fingerprint.Model, "default"), "transport errors cannot clear the latest valid vision result")
			require.NoError(t, db.Model(&failedVision).Update("quality", "degraded").Error)
			require.NoError(t, db.Model(&vision).Updates(map[string]any{"quality": "normal", "result": `{"closest":"这是音箱"}`}).Error)
			InvalidateGroupQualityCache()
			assert.False(t, requestGroupDegraded(ctx(), fingerprint.Model, "default"))
			recent, err = RecentDegradedGroups(ctx().Request.Context())
			require.NoError(t, err)
			assert.False(t, recent["default"])
			require.NoError(t, db.Model(&vision).Updates(map[string]any{"quality": "degraded", "created_at": now - 8*3600, "result": `{"closest":"石头"}`}).Error)
			InvalidateGroupQualityCache()
			assert.True(t, requestGroupDegraded(ctx(), fingerprint.Model, "default"), "routing keeps the latest valid result")
			recent, err = RecentDegradedGroups(ctx().Request.Context())
			require.NoError(t, err)
			assert.False(t, recent["default"], "recent tag expires after the existing six-hour window")
			legacyAnswer := `{"closest":"我这边没收到可查看的图片或文件，暂时无法判断是什么。请重新上传图片"}`
			require.NoError(t, db.Model(&vision).Updates(map[string]any{"result": legacyAnswer, "status": "failed", "error": "vision_keyword_missing"}).Error)
			InvalidateGroupQualityCache()
			unavailable, err := VisionUnavailableGroups(ctx().Request.Context())
			require.NoError(t, err)
			assert.True(t, unavailable["default"], "legacy answers remain effective beyond the recent window and later transport failures")
			assert.False(t, requestGroupDegraded(ctx(), fingerprint.Model, "default"))
			recent, err = RecentDegradedGroups(ctx().Request.Context())
			require.NoError(t, err)
			assert.False(t, recent["default"])
			require.NoError(t, db.Model(&vision).Update("created_at", now-2).Error)
			summaries, err := GroupDetectionSummaries(ctx().Request.Context(), time.Unix(now, 0), nil, nil)
			require.NoError(t, err)
			summary := summaries["default"].Vision
			require.NotNil(t, summary)
			assert.Equal(t, 1, summary.Incomplete)
			assert.Equal(t, "incomplete", summary.LatestValid.Quality)
			var unchanged model.DetectionRun
			require.NoError(t, db.First(&unchanged, "id = ?", vision.ID).Error)
			assert.Equal(t, "degraded", unchanged.Quality, "reading history must not rewrite stored results")
			ciallo := model.DetectionRun{ID: "ciallo-degraded", RequestKey: "ciallo-degraded", ChannelID: 1, Kind: "ciallo", CreatedAt: now, Settled: true, Status: "succeeded", Quality: "degraded"}
			require.NoError(t, db.Create(&ciallo).Error)
			InvalidateGroupQualityCache()
			assert.True(t, requestGroupDegraded(ctx(), fingerprint.Model, "default"), "incomplete vision must not clear Ciallo degradation")
			require.NoError(t, db.Model(&vision).Updates(map[string]any{"status": "succeeded", "result": `{"closest":"这是音箱"}`, "created_at": now}).Error)
			InvalidateGroupQualityCache()
			unavailable, err = VisionUnavailableGroups(ctx().Request.Context())
			require.NoError(t, err)
			assert.False(t, unavailable["default"])
			assert.True(t, requestGroupDegraded(ctx(), fingerprint.Model, "default"), "vision recovery must not clear Ciallo degradation")
			require.NoError(t, db.Model(&ciallo).Update("quality", "normal").Error)
			InvalidateGroupQualityCache()
			assert.False(t, requestGroupDegraded(ctx(), fingerprint.Model, "default"))
			require.NoError(t, db.Delete(&manual).Error)
			require.NoError(t, model.SaveStatusPageGroup(ctx().Request.Context(), "default", nil))
			groupDegradationSnapshot.updated = time.Time{}
			assert.False(t, requestGroupDegraded(ctx(), fingerprint.Model, "default"))
			unavailable, err = VisionUnavailableGroups(ctx().Request.Context())
			require.NoError(t, err)
			assert.False(t, unavailable["default"], "administrator reset removes the vision conclusion")
		})
	}
}

func TestGroupDegradationRetainsFullHistoryForNotifications(t *testing.T) {
	now := time.Now()
	channels := []model.Channel{{Id: 1, Group: "default", Models: fingerprint.Model, Status: 1}}
	rounds := []AstraSamplingRound{{ID: "valid", ChannelID: 1, StartedAt: now.Unix() - 100, FinishedAt: now.Unix() - 90, Result: fingerprint.Result{Engine: fingerprint.Version, Status: "mismatch", Closest: "gpt-5.5", ValidSamples: 1}}}
	for i := range 15 {
		rounds = append(rounds, AstraSamplingRound{ChannelID: 1, StartedAt: now.Unix() - int64(i), FinishedAt: now.Unix(), Result: fingerprint.Result{Engine: fingerprint.Version, Status: "error"}})
	}
	body, err := common.Marshal(AstraSamplingState{Rounds: rounds})
	require.NoError(t, err)
	tasks := []*model.SystemTask{{Status: model.SystemTaskStatusSucceeded, State: string(body)}}
	summary := AggregateAstraSampling(channels, tasks, now)["default"]
	require.NotNil(t, summary.LatestValid)
	assert.Equal(t, "gpt-5.5", summary.LatestValid.Closest)
	assert.Len(t, summary.History, 16)
	rounds = append(rounds, AstraSamplingRound{ID: "recovered", ChannelID: 1, StartedAt: now.Unix(), FinishedAt: now.Unix(), Result: fingerprint.Result{Engine: fingerprint.Version, Status: "match", Closest: fingerprint.Model, ValidSamples: 2}})
	body, err = common.Marshal(AstraSamplingState{Rounds: rounds})
	require.NoError(t, err)
	tasks[0].State = string(body)
	summary = AggregateAstraSampling(channels, tasks, now)["default"]
	require.NotNil(t, summary.LatestValid)
	assert.Equal(t, fingerprint.Model, summary.LatestValid.Closest)
	reset := aggregateAstraSampling(channels, tasks, now, map[string]model.StatusPageGroup{"default": {ResetBefore: now.Unix()}})
	assert.Nil(t, reset["default"].LatestValid)
}

func TestGroupDegradationOverridesExistingPoolAffinity(t *testing.T) {
	db := setupChannelSelectAutoGroupsTest(t)
	createChannelSelectAutoGroupsChannel(t, db, 6401, "default", fingerprint.Model)
	createChannelSelectAutoGroupsChannel(t, db, 6402, "vip", fingerprint.Model)
	model.InitChannelCache()
	ResetGroupAvailability()
	t.Cleanup(ResetGroupAvailability)
	previousRedis := common.RedisEnabled
	common.RedisEnabled = false
	t.Cleanup(func() { common.RedisEnabled = previousRedis; _ = poolBindings().Purge() })
	require.NoError(t, poolBindings().Purge())
	config := operation_setting.GetChannelAffinitySetting()
	previous := *config
	*config = operation_setting.ChannelAffinitySetting{Enabled: true, DefaultTTLSeconds: 3600, SessionMode: "strict", Rules: []operation_setting.ChannelAffinityRule{{Name: "degradation-test", ModelRegex: []string{".*"}, KeySources: []operation_setting.ChannelAffinityKeySource{{Type: "request_header", Key: "Session-Id"}}, SessionMode: "strict"}}}
	t.Cleanup(func() { *config = previous })
	ctx := func(degraded bool) *gin.Context {
		c := setupGroupPoolContext(t)
		c.Request = httptest.NewRequest("POST", "/v1/responses", nil)
		c.Request.Header.Set("Session-Id", "degradation-session")
		c.Set("id", 100)
		c.Set("token_id", 100)
		c.Set(model.TokenRoutingContextKey, &model.TokenRoutingPolicy{Version: 1, Mode: "custom", SkipDegradedGroups: true, Groups: []model.TokenGroupEntry{{Group: "default", Priority: 0}, {Group: "vip", Priority: 1}}})
		rounds := map[string]AstraSamplingHistory{}
		if degraded {
			rounds["default"] = AstraSamplingHistory{At: time.Now().Unix(), Closest: "gpt-5.5", ValidSamples: 1}
		}
		c.Set("pool_group_fingerprints", rounds)
		return c
	}
	first := ctx(false)
	selected, group, err := SelectChannelForRequest(first, fingerprint.Model, &RetryParam{Ctx: first, TokenGroup: "default", ModelName: fingerprint.Model})
	require.Nil(t, err)
	require.NotNil(t, selected)
	require.Equal(t, 6401, selected.Id)
	first.Set("channel_id", selected.Id)
	common.SetContextKey(first, constant.ContextKeyAutoGroup, group)
	RequestPolicy(first).Successful = true
	RecordPoolAffinity(first)
	bound, _ := PreferredPoolChannel(ctx(false), fingerprint.Model)
	require.NotNil(t, bound)
	image := ctx(false)
	image.Request = httptest.NewRequest("POST", "/v1/responses", strings.NewReader(`{"input":[{"role":"user","content":[{"type":"input_image","image_url":"data:image/png;base64,eA=="}]}]}`))
	image.Request.Header.Set("Session-Id", "degradation-session")
	image.Set("pool_group_visions", map[string]DetectionHistory{"default": {Quality: "incomplete"}})
	selected, group, err = SelectChannelForRequest(image, fingerprint.Model, &RetryParam{Ctx: image, TokenGroup: "default", ModelName: fingerprint.Model})
	require.Nil(t, err)
	require.NotNil(t, selected)
	assert.Equal(t, 6401, selected.Id, "incomplete vision must preserve the existing session affinity")
	assert.Equal(t, "default", group)
	next := ctx(true)
	selected, group, err = SelectChannelForRequest(next, fingerprint.Model, &RetryParam{Ctx: next, TokenGroup: "default", ModelName: fingerprint.Model})
	require.Nil(t, err)
	require.NotNil(t, selected)
	assert.Equal(t, 6402, selected.Id)
	assert.Equal(t, "vip", group)
}

func TestVisionIncompleteRoutingAcrossModelsAndPolicies(t *testing.T) {
	db := setupChannelSelectAutoGroupsTest(t)
	for i, group := range []string{"default", "vip"} {
		createChannelSelectAutoGroupsChannel(t, db, 6601+i, group, fingerprint.Model)
		require.NoError(t, db.Model(&model.Channel{}).Where("id = ?", 6601+i).Update("models", fingerprint.Model+",gpt-6-sol").Error)
		require.NoError(t, db.Create(&model.Ability{Group: group, Model: "gpt-6-sol", ChannelId: 6601 + i, Enabled: true, Weight: 100}).Error)
	}
	model.InitChannelCache()
	ResetGroupAvailability()
	t.Cleanup(ResetGroupAvailability)
	for _, mode := range []string{"smart", "custom"} {
		for _, allow := range []bool{false, true} {
			for _, modelName := range []string{fingerprint.Model, "gpt-6-sol"} {
				for _, image := range []bool{false, true} {
					c := setupGroupPoolContext(t)
					c.Request = httptest.NewRequest("POST", "/v1/responses", nil)
					c.Set(model.TokenRoutingContextKey, &model.TokenRoutingPolicy{Version: 1, Mode: mode, AllowDegraded: allow, SkipDegradedGroups: !allow, Groups: []model.TokenGroupEntry{{Group: "default"}, {Group: "vip", Priority: 1}}})
					if image {
						c.Request = httptest.NewRequest("POST", "/v1/responses", strings.NewReader(`{"input":[{"role":"user","content":[{"type":"input_image","image_url":"data:image/png;base64,eA=="}]}]}`))
					}
					c.Set("pool_group_fingerprints", map[string]AstraSamplingHistory{})
					visions := map[string]DetectionHistory{"default": {Quality: "incomplete"}, "vip": {Quality: "incomplete"}}
					c.Set("pool_group_visions", visions)
					channel, _, err := selectTokenRoutingPool(&RetryParam{Ctx: c, TokenGroup: "default", ModelName: modelName})
					require.NoError(t, err)
					require.NotNil(t, channel)
					assert.Equal(t, 6601, channel.Id, "mode=%s allow=%v model=%s image=%v", mode, allow, modelName, image)
					assert.False(t, poolGroupSkipped(c, modelName, "default"), "incomplete vision must not exclude retries or cooldown recovery")
				}
			}
		}
	}
}

func TestGroupDegradationSkipsPoolGroupsAndRespectsUnknownRecoveryAndOptOut(t *testing.T) {
	db := setupChannelSelectAutoGroupsTest(t)
	createChannelSelectAutoGroupsChannel(t, db, 6301, "default", fingerprint.Model)
	createChannelSelectAutoGroupsChannel(t, db, 6302, "vip", fingerprint.Model)
	model.InitChannelCache()
	ResetGroupAvailability()
	t.Cleanup(ResetGroupAvailability)
	for _, test := range []struct {
		name          string
		enabled       bool
		closest       string
		age           time.Duration
		degradeBackup bool
		want          int
	}{
		{"latest mismatch switches", true, "gpt-5.5", time.Minute, false, 6302},
		{"recovered group returns", true, fingerprint.Model, time.Minute, false, 6301},
		{"unknown stays eligible", true, "", time.Minute, false, 6301},
		{"valid result persists beyond window", true, "gpt-5.5", 7 * time.Hour, false, 6302},
		{"disabled keeps routing", false, "gpt-5.5", time.Minute, false, 6301},
		{"all degraded has no route", true, "gpt-5.5", time.Minute, true, 0},
	} {
		t.Run(test.name, func(t *testing.T) {
			c := setupGroupPoolContext(t)
			c.Set(model.TokenRoutingContextKey, &model.TokenRoutingPolicy{Version: 1, Mode: "custom", SkipDegradedGroups: test.enabled, Groups: []model.TokenGroupEntry{{Group: "default", Priority: 0}, {Group: "vip", Priority: 1}}})
			rounds := map[string]AstraSamplingHistory{"default": {At: time.Now().Add(-test.age).Unix(), Closest: test.closest, ValidSamples: 1}}
			if test.degradeBackup {
				rounds["vip"] = rounds["default"]
			}
			c.Set("pool_group_fingerprints", rounds)
			selected, _, err := selectTokenRoutingPool(&RetryParam{Ctx: c, TokenGroup: "default", ModelName: fingerprint.Model})
			require.NoError(t, err)
			if test.want == 0 {
				assert.Nil(t, selected)
				return
			}
			require.NotNil(t, selected)
			assert.Equal(t, test.want, selected.Id)
		})
	}
}