package model

import (
	"fmt"
	"github.com/QuantumNous/new-api/common"
	"github.com/glebarez/sqlite"
	"github.com/stretchr/testify/require"
	"gorm.io/driver/mysql"
	"gorm.io/driver/postgres"
	"gorm.io/gorm"
	"gorm.io/gorm/schema"
	"os"
	"path/filepath"
	"sync"
	"testing"
	"time"
)

func TestTokenRoutingResolution(t *testing.T) {
	p := TokenRoutingPolicy{Version: 1, Mode: "custom", Groups: []TokenGroupEntry{{"a", 0}, {"b", 0}, {"c", 2}}, ModelRules: []ModelRoutingRule{{Model: "model-x", Groups: []TokenGroupEntry{{"c", 0}, {"b", 0}}}}}
	allowed := func(string) bool { return true }
	require.NoError(t, p.Validate(6, allowed))
	require.Equal(t, []TokenGroupEntry{{"c", 0}, {"b", 0}, {"a", 1}}, p.Resolve("model-x", nil, allowed))
	require.Equal(t, []TokenGroupEntry{{"a", 0}, {"b", 0}, {"c", 1}}, p.Resolve("MODEL-X", nil, allowed))
}

func TestTokenRoutingLegacy(t *testing.T) {
	for _, tc := range []struct {
		token  Token
		mode   string
		groups []TokenGroupEntry
	}{
		{Token{Group: "auto"}, "inherit", []TokenGroupEntry{}},
		{Token{Group: ""}, "user", []TokenGroupEntry{}},
		{Token{Group: "a", AutoGroups: `["b","c"]`, CrossGroupRetry: true}, "custom", []TokenGroupEntry{{"a", 0}, {"b", 1}, {"c", 2}}},
		{Token{Group: "a", AutoGroups: `["b"]`, CrossGroupRetry: false}, "custom", []TokenGroupEntry{{"a", 0}}},
		{Token{Group: "a", AutoGroups: `[{"group":"b","priority":0},{"group":"a","priority":0}]`, CrossGroupRetry: true}, "custom", []TokenGroupEntry{{"b", 0}, {"a", 0}}},
	} {
		p, err := tc.token.LegacyRoutingPolicy()
		require.NoError(t, err)
		require.Equal(t, tc.mode, p.Mode)
		require.Equal(t, tc.groups, p.Groups)
		require.Empty(t, p.ModelRules)
	}
	p := TokenRoutingPolicy{Version: 1, Mode: "custom", Groups: []TokenGroupEntry{{"a", 0}}, ModelRules: []ModelRoutingRule{{Model: "x", Groups: []TokenGroupEntry{{"b", 0}}}}}
	require.Error(t, p.Validate(6, func(string) bool { return true }))
}

func TestRelayAttemptDatabaseMatrix(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":
				dsn := os.Getenv("TEST_MYSQL_DSN")
				if dsn == "" {
					t.Skip("TEST_MYSQL_DSN not configured")
				}
				driver = mysql.Open(dsn)
			case "postgres":
				dsn := os.Getenv("TEST_POSTGRES_DSN")
				if dsn == "" {
					t.Skip("TEST_POSTGRES_DSN not configured")
				}
				driver = postgres.Open(dsn)
			}
			db, err := gorm.Open(driver, &gorm.Config{NamingStrategy: schema.NamingStrategy{TablePrefix: fmt.Sprintf("attempt_%d_", time.Now().UnixNano())}})
			require.NoError(t, err)
			sqlDB, err := db.DB()
			require.NoError(t, err)
			sqlDB.SetMaxOpenConns(1)
			oldDB, oldLog, oldType, oldRedis := DB, LOG_DB, common.MainDatabaseType(), common.RedisEnabled
			DB = db
			LOG_DB = db
			common.SetMainDatabaseType(common.DatabaseType(dialect))
			common.RedisEnabled = false
			initCol()
			models := []any{&User{}, &Token{}, &Channel{}, &RelayAttempt{}, &Log{}, &QuotaData{}, &SubscriptionPlan{}, &UserSubscription{}}
			t.Cleanup(func() {
				for _, m := range models {
					require.NoError(t, db.Migrator().DropTable(m))
				}
				DB = oldDB
				LOG_DB = oldLog
				common.SetMainDatabaseType(oldType)
				common.RedisEnabled = oldRedis
				initCol()
				require.NoError(t, sqlDB.Close())
			})
			for range 2 {
				require.NoError(t, db.AutoMigrate(models...))
			}
			user := User{Username: "attempt-owner", Quota: 10000}
			require.NoError(t, db.Create(&user).Error)
			token := Token{UserId: user.Id, Key: "attempt-key", Status: common.TokenStatusEnabled, RemainQuota: 10000, Group: "a", CrossGroupRetry: true, AutoGroups: `["b"]`}
			require.NoError(t, db.Create(&token).Error)
			channel := Channel{Name: "attempt-channel", Status: common.ChannelStatusEnabled}
			require.NoError(t, db.Create(&channel).Error)
			// Upgrade a populated legacy token table, then run startup migrations twice.
			require.NoError(t, db.Migrator().DropColumn(&Token{}, "RoutingPolicy"))
			require.NoError(t, db.AutoMigrate(&Token{}))
			count, err := MigrateTokenRoutingPolicies()
			require.NoError(t, err)
			require.Equal(t, 1, count)
			count, err = MigrateTokenRoutingPolicies()
			require.NoError(t, err)
			require.Zero(t, count)
			var migrated Token
			require.NoError(t, db.First(&migrated, token.Id).Error)
			p, err := migrated.GetRoutingPolicy()
			require.NoError(t, err)
			require.Equal(t, []TokenGroupEntry{{"a", 0}, {"b", 1}}, p.Groups)
			var total int
			for i, quota := range []int{120, 230, 0} {
				a := RelayAttempt{AttemptKey: fmt.Sprintf("request:%d", i), RequestID: "request", Number: i + 1, UserID: user.Id, TokenID: token.Id, ChannelID: channel.Id, Reserved: 300}
				require.NoError(t, ReserveRelayAttempt(&a, "wallet_only"))
				require.NoError(t, ReserveRelayAttempt(&a, "wallet_only"))
				require.NoError(t, CheckpointRelayAttempt(a.ID, quota, `{"total_tokens":10}`, "", "luna", "model_mismatch", "upstream", quota > 0, time.Now().Unix()))
				log := &Log{UserId: user.Id, Type: LogTypeConsume, RequestId: "request"}
				require.NoError(t, SettleRelayAttempt(a.ID, quota, log))
				require.NoError(t, SettleRelayAttempt(a.ID, quota+1, log))
				total += quota
			}
			require.NoError(t, DeliverRelayAttemptLogs())
			require.NoError(t, DeliverRelayAttemptLogs())
			require.NoError(t, db.First(&user, user.Id).Error)
			require.Equal(t, 10000-total, user.Quota)
			require.Equal(t, total, user.UsedQuota)
			require.Equal(t, 2, user.RequestCount)
			require.NoError(t, db.First(&token, token.Id).Error)
			require.Equal(t, 10000-total, token.RemainQuota)
			require.Equal(t, total, token.UsedQuota)
			require.NoError(t, db.First(&channel, channel.Id).Error)
			require.EqualValues(t, total, channel.UsedQuota)
			var logs int64
			require.NoError(t, db.Model(&Log{}).Count(&logs).Error)
			require.EqualValues(t, 3, logs)
			// A log-store outage never rolls back an already-committed financial fact.
			require.NoError(t, db.Model(&RelayAttempt{}).Where("request_id = ?", "request").Update("log_delivered", false).Error)
			require.NoError(t, DeliverRelayAttemptLogs())
			require.NoError(t, db.Model(&Log{}).Count(&logs).Error)
			require.EqualValues(t, 3, logs)
			plan := SubscriptionPlan{Title: "attempt-sub", QuotaResetPeriod: "never"}
			require.NoError(t, db.Create(&plan).Error)
			sub := UserSubscription{UserId: user.Id, PlanId: plan.Id, AmountTotal: 1000, Status: "active", EndTime: time.Now().Add(time.Hour).Unix()}
			require.NoError(t, db.Create(&sub).Error)
			subAttempt := RelayAttempt{AttemptKey: "sub-request", UserID: user.Id, TokenID: token.Id, ChannelID: channel.Id, Reserved: 300}
			require.NoError(t, ReserveRelayAttempt(&subAttempt, "subscription_only"))
			require.Equal(t, "subscription", subAttempt.Funding)
			require.NoError(t, SettleRelayAttempt(subAttempt.ID, 120, nil))
			require.NoError(t, SettleRelayAttempt(subAttempt.ID, 999, nil))
			require.NoError(t, db.First(&sub, sub.Id).Error)
			require.EqualValues(t, 120, sub.AmountUsed)
			require.NoError(t, db.First(&user, user.Id).Error)
			require.Equal(t, 10000-total, user.Quota)
			rejected := RelayAttempt{AttemptKey: "too-expensive", UserID: user.Id, TokenID: token.Id, Reserved: 20000}
			require.ErrorIs(t, ReserveRelayAttempt(&rejected, "wallet_only"), ErrAttemptQuota)
			require.NoError(t, db.First(&user, user.Id).Error)
			require.Equal(t, 10000-total, user.Quota)
		})
	}
}

func TestRelayAttemptSQLiteConcurrentSettlementAndReserve(t *testing.T) {
	db, err := gorm.Open(sqlite.Open(filepath.Join(t.TempDir(), "attempt.db")), &gorm.Config{})
	require.NoError(t, err)
	sqlDB, err := db.DB()
	require.NoError(t, err)
	sqlDB.SetMaxOpenConns(8)
	oldDB, oldType, oldRedis := DB, common.MainDatabaseType(), common.RedisEnabled
	DB = db
	common.SetMainDatabaseType(common.DatabaseType("sqlite"))
	common.RedisEnabled = false
	initCol()
	t.Cleanup(func() {
		DB = oldDB
		common.SetMainDatabaseType(oldType)
		common.RedisEnabled = oldRedis
		initCol()
		_ = sqlDB.Close()
	})
	require.NoError(t, db.AutoMigrate(&User{}, &Token{}, &Channel{}, &RelayAttempt{}))
	user := User{Username: "concurrent", Quota: 10000}
	require.NoError(t, db.Create(&user).Error)
	token := Token{UserId: user.Id, Key: "concurrent-token", Status: common.TokenStatusEnabled, RemainQuota: 10000}
	require.NoError(t, db.Create(&token).Error)
	channel := Channel{Name: "concurrent"}
	require.NoError(t, db.Create(&channel).Error)
	var wg sync.WaitGroup
	failures := make(chan error, 16)
	for i := 0; i < 16; i++ {
		wg.Add(1)
		go func(i int) {
			defer wg.Done()
			a := RelayAttempt{AttemptKey: fmt.Sprintf("concurrent:%d", i), UserID: user.Id, TokenID: token.Id, ChannelID: channel.Id, Reserved: 100}
			if err := ReserveRelayAttempt(&a, "wallet_only"); err != nil {
				failures <- err
				return
			}
			if err := SettleRelayAttempt(a.ID, 50, nil); err != nil {
				failures <- err
				return
			}
			if err := SettleRelayAttempt(a.ID, 50, nil); err != nil {
				failures <- err
			}
		}(i)
	}
	wg.Wait()
	close(failures)
	for err := range failures {
		require.NoError(t, err)
	}
	require.NoError(t, db.First(&user, user.Id).Error)
	require.Equal(t, 9200, user.Quota)
	require.Equal(t, 800, user.UsedQuota)
	require.NoError(t, db.First(&token, token.Id).Error)
	require.Equal(t, 9200, token.RemainQuota)
	require.Equal(t, 800, token.UsedQuota)
	var count int64
	require.NoError(t, db.Model(&RelayAttempt{}).Where("state = ?", "settled").Count(&count).Error)
	require.EqualValues(t, 16, count)
}