package service

import (
	"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"
	"github.com/QuantumNous/new-api/setting"
	"github.com/QuantumNous/new-api/setting/model_setting"
	"github.com/QuantumNous/new-api/setting/ratio_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"
	"os"
	"testing"
	"time"
)

func TestAutomaticPoolModelsDatabaseMatrix(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("no isolated MySQL database")
				}
				driver = mysql.Open(os.Getenv("TEST_MYSQL_DSN"))
			case "postgres":
				if os.Getenv("TEST_POSTGRES_DSN") == "" {
					t.Skip("no isolated PostgreSQL database")
				}
				driver = postgres.Open(os.Getenv("TEST_POSTGRES_DSN"))
			}
			db, err := gorm.Open(driver, &gorm.Config{NamingStrategy: schema.NamingStrategy{TablePrefix: fmt.Sprintf("automodel_%d_", time.Now().UnixNano())}})
			require.NoError(t, err)
			oldDB, oldRedis, oldMemory := model.DB, common.RedisEnabled, common.MemoryCacheEnabled
			oldGroups, oldRatios := setting.UserUsableGroups2JSONString(), ratio_setting.GroupRatio2JSONString()
			model.DB = db
			common.RedisEnabled = false
			common.MemoryCacheEnabled = false
			require.NoError(t, db.AutoMigrate(&model.Channel{}, &model.Ability{}, &model.Token{}))
			t.Cleanup(func() {
				require.NoError(t, db.Migrator().DropTable(&model.Token{}, &model.Ability{}, &model.Channel{}))
				model.DB = oldDB
				common.RedisEnabled = oldRedis
				common.MemoryCacheEnabled = oldMemory
				require.NoError(t, setting.UpdateUserUsableGroupsByJSONString(oldGroups))
				require.NoError(t, ratio_setting.UpdateGroupRatioByJSONString(oldRatios))
				conn, e := db.DB()
				require.NoError(t, e)
				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)
			require.NoError(t, setting.UpdateUserUsableGroupsByJSONString(`{"auto-a":"","auto-b":""}`))
			require.NoError(t, ratio_setting.UpdateGroupRatioByJSONString(`{"auto-a":1,"auto-b":1,"private":1}`))
			require.NoError(t, db.Create(&[]model.Channel{{Id: 561, Status: 1}, {Id: 562, Status: 1}, {Id: 563, Status: 2}, {Id: 564, Status: 1}}).Error)
			require.NoError(t, db.Create(&[]model.Ability{
				{ChannelId: 561, Group: "auto-a", Model: "model-a", Enabled: true},
				{ChannelId: 561, Group: "auto-a", Model: "shared", Enabled: true},
				{ChannelId: 562, Group: "auto-b", Model: "shared", Enabled: true},
				{ChannelId: 562, Group: "auto-b", Model: "gpt-6-astra", Enabled: true},
				{ChannelId: 563, Group: "auto-b", Model: "disabled-channel", Enabled: true},
				{ChannelId: 562, Group: "auto-b", Model: "disabled-ability", Enabled: false},
				{ChannelId: 564, Group: "private", Model: "unauthorized", Enabled: true},
			}).Error)
			for _, preference := range []string{"cost", "speed", "stable"} {
				c, _ := gin.CreateTestContext(httptest.NewRecorder())
				c.Request = httptest.NewRequest("GET", "/v1/models", nil)
				common.SetContextKey(c, constant.ContextKeyUserGroup, "auto-a")
				c.Set(model.TokenRoutingContextKey, &model.TokenRoutingPolicy{Version: 1, Mode: "smart", AutoPreference: preference})
				c.Set("pool_failures", map[string]int{"auto-a": 2, "auto-b": 2})
				require.Equal(t, []string{"gpt-6-astra", "model-a", "shared"}, AutomaticPoolModels(c))
			}
			require.NoError(t, db.Create(&model.Channel{Id: 565, Status: 1}).Error)
			require.NoError(t, db.Create(&model.Ability{ChannelId: 565, Group: "new-group", Model: "new-model", Enabled: true}).Error)
			require.NoError(t, setting.UpdateUserUsableGroupsByJSONString(`{"auto-a":"","auto-b":"","new-group":""}`))
			require.NoError(t, ratio_setting.UpdateGroupRatioByJSONString(`{"auto-a":1,"auto-b":1,"new-group":1}`))
			c, _ := gin.CreateTestContext(httptest.NewRecorder())
			c.Request = httptest.NewRequest("GET", "/v1/models", nil)
			common.SetContextKey(c, constant.ContextKeyUserGroup, "auto-a")
			require.Equal(t, []string{"gpt-6-astra", "model-a", "new-model", "shared"}, AutomaticPoolModels(c))

			cfg := model_setting.GetGlobalSettings()
			oldImageGroup, oldImageModels := cfg.ImageGenerationGroup, cfg.ImageGenerationModels
			cfg.ImageGenerationGroup, cfg.ImageGenerationModels = "auto-b", []string{"blacklist-image"}
			t.Cleanup(func() {
				cfg.ImageGenerationGroup, cfg.ImageGenerationModels = oldImageGroup, oldImageModels
				automaticMetricCache.Purge()
				_ = poolCooldowns().Purge()
			})
			createChannelSelectAutoGroupsChannel(t, db, 571, "auto-a", "blacklist-model")
			createChannelSelectAutoGroupsChannel(t, db, 572, "auto-b", "blacklist-model")
			createChannelSelectAutoGroupsChannel(t, db, 573, "auto-a", "blacklist-image")
			createChannelSelectAutoGroupsChannel(t, db, 574, "auto-b", "blacklist-image")
			common.MemoryCacheEnabled = true
			model.InitChannelCache()
			automaticMetricCache.Set("blacklist-model", perfmetrics.QueryResult{})
			automaticMetricCache.Set("blacklist-image", perfmetrics.QueryResult{})
			makeContext := func(policy *model.TokenRoutingPolicy, tokenID int) *gin.Context {
				ctx, _ := gin.CreateTestContext(httptest.NewRecorder())
				ctx.Request = httptest.NewRequest("POST", "/v1/responses", nil)
				ctx.Set("token_id", tokenID)
				ctx.Set("id", 7)
				common.SetContextKey(ctx, constant.ContextKeyUserGroup, "auto-a")
				ctx.Set(model.TokenRoutingContextKey, policy)
				return ctx
			}
			for _, preference := range []string{"cost", "speed", "stable"} {
				for _, blocked := range [][]string{nil, {"auto-a"}, {"auto-a", "auto-b"}} {
					policy := &model.TokenRoutingPolicy{Version: 1, Mode: "smart", AutoPreference: preference, ExcludedGroups: blocked}
					body, err := common.Marshal(policy)
					require.NoError(t, err)
					token := model.Token{Key: fmt.Sprintf("blacklist-%s-%d", preference, len(blocked)), UserId: 7, RoutingPolicy: string(body)}
					require.NoError(t, token.Insert())
					var stored model.Token
					require.NoError(t, db.First(&stored, token.Id).Error)
					loaded, err := stored.GetRoutingPolicy()
					require.NoError(t, err)
					assert.Equal(t, blocked, loaded.ExcludedGroups)
					ctx := makeContext(loaded, stored.Id)
					pool := ResolveRequestPool(ctx, "blacklist-model")
					require.Len(t, pool, 2-len(blocked))
					assert.Equal(t, max(0, 2*len(pool)-1), PoolRetryBudget(ctx, "blacklist-model"))
					for _, entry := range pool {
						assert.False(t, loaded.ExcludesGroup(entry.Group))
					}
					channel, group, selectErr := SelectChannelForRequest(ctx, "blacklist-model", &RetryParam{Ctx: ctx, TokenGroup: "auto", ModelName: "blacklist-model"})
					if len(pool) == 0 {
						assert.Nil(t, channel)
						require.NotNil(t, selectErr)
						assert.True(t, selectErr.NoAvailableChannel)
					} else {
						require.Nil(t, selectErr)
						require.NotNil(t, channel)
						assert.False(t, loaded.ExcludesGroup(group))
						RecordPoolRetryableFailure(ctx, "blacklist-model", group)
						retryChannel, retryGroup, err := selectTokenRoutingPool(&RetryParam{Ctx: ctx, TokenGroup: "auto", ModelName: "blacklist-model", Retry: common.GetPointer(1)})
						require.NoError(t, err)
						require.NotNil(t, retryChannel)
						assert.False(t, loaded.ExcludesGroup(retryGroup))
					}
					loaded.ExcludedGroups = nil
					body, err = common.Marshal(loaded)
					require.NoError(t, err)
					stored.RoutingPolicy = string(body)
					require.NoError(t, stored.Update())
					require.NoError(t, db.First(&stored, token.Id).Error)
					cleared, err := stored.GetRoutingPolicy()
					require.NoError(t, err)
					assert.Empty(t, cleared.ExcludedGroups)
					assert.Len(t, ResolveRequestPool(makeContext(cleared, stored.Id), "blacklist-model"), 2)
				}
			}
			policy := &model.TokenRoutingPolicy{Version: 1, Mode: "smart", ExcludedGroups: []string{"auto-b"}}
			ctx := makeContext(policy, 900)
			assert.Empty(t, ResolveRequestPool(ctx, "blacklist-image"), "the dedicated image group cannot bypass the blacklist or fall back")
			assert.NotContains(t, AutomaticPoolModels(ctx), "gpt-6-astra")
			assert.Contains(t, AutomaticPoolModels(ctx), "shared")
			assert.Contains(t, AutomaticPoolModels(makeContext(&model.TokenRoutingPolicy{Version: 1, Mode: "smart"}, 901)), "gpt-6-astra", "another key is unaffected")
			manual := &model.TokenRoutingPolicy{Version: 1, Mode: "custom", Groups: []model.TokenGroupEntry{{Group: "auto-b", Priority: 0}}, ExcludedGroups: []string{"auto-b"}}
			assert.Equal(t, manual.Groups, ResolveRequestPool(makeContext(manual, 902), "blacklist-model"))
		})
	}
}