package service

import (
	"fmt"
	"github.com/QuantumNous/new-api/common"
	"github.com/QuantumNous/new-api/constant"
	"github.com/QuantumNous/new-api/model"
	"github.com/QuantumNous/new-api/setting"
	"github.com/QuantumNous/new-api/setting/ratio_setting"
	"github.com/gin-gonic/gin"
	"github.com/glebarez/sqlite"
	"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{}))
			t.Cleanup(func() {
				require.NoError(t, db.Migrator().DropTable(&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))
		})
	}
}