package service

import (
	"fmt"
	"os"
	"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
				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"))
			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"))
		})
	}
}

func TestGroupDegradationUsesLatestValidBeforeHistoryTruncation(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, 12)
	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)
	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 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)
		})
	}
}