package model

import (
	"context"
	"fmt"
	"github.com/QuantumNous/new-api/common"
	"os"
	"testing"
	"time"

	"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"
)

func TestGroupModelCostsDatabaseMatrix(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 MySQL test database")
				}
				driver = mysql.Open(os.Getenv("TEST_MYSQL_DSN"))
			case "postgres":
				if os.Getenv("TEST_POSTGRES_DSN") == "" {
					t.Skip("no PostgreSQL test database")
				}
				driver = postgres.Open(os.Getenv("TEST_POSTGRES_DSN"))
			}
			db, err := gorm.Open(driver, &gorm.Config{NamingStrategy: schema.NamingStrategy{TablePrefix: fmt.Sprintf("cost_%d_", time.Now().UnixNano())}})
			require.NoError(t, err)
			oldDB, oldUnit := LOG_DB, common.QuotaPerUnit
			LOG_DB = db
			common.QuotaPerUnit = 500000
			t.Cleanup(func() {
				LOG_DB = oldDB
				common.QuotaPerUnit = oldUnit
				require.NoError(t, db.Migrator().DropTable(&Log{}))
				sqlDB, e := db.DB()
				require.NoError(t, e)
				require.NoError(t, sqlDB.Close())
			})
			require.NoError(t, db.AutoMigrate(&Log{}))
			logs := []Log{
				{Type: 2, TokenId: 1, CreatedAt: 150, Group: "a", ModelName: "m", Quota: 50, PromptTokens: 100, CompletionTokens: 100, Other: `{"cache_tokens":80}`},
				{Type: 2, TokenId: 2, CreatedAt: 150, Group: "a", ModelName: "m", Quota: 150, PromptTokens: 200, CompletionTokens: 100, Other: `{"usage_semantic":"anthropic","cache_tokens":100,"cache_creation_tokens":400}`},
				{Type: 2, TokenId: 3, CreatedAt: 150, Group: "b", ModelName: "m", Quota: 1000, PromptTokens: 200, CompletionTokens: 100, Other: `{"input_tokens_total":900}`},
				{Type: 2, TokenId: 1, CreatedAt: 150, Group: "a", ModelName: "free", Quota: 0, PromptTokens: 100},
				{Type: 2, TokenId: 0, CreatedAt: 150, Group: "a", ModelName: "m", Quota: 999999, PromptTokens: 100},
				{Type: 2, TokenId: 1, CreatedAt: 99, Group: "a", ModelName: "m", Quota: 999999, PromptTokens: 100},
				{Type: 2, TokenId: 1, CreatedAt: 200, Group: "a", ModelName: "m", Quota: 999999, PromptTokens: 100},
				{Type: 5, TokenId: 1, CreatedAt: 150, Group: "a", ModelName: "m", Quota: 999999, PromptTokens: 100},
				{Type: 2, TokenId: 1, CreatedAt: 150, Group: "a", ModelName: "no-tokens", Quota: 500},
			}
			require.NoError(t, db.Create(&logs).Error)
			result, err := GetGroupModelCosts(context.Background(), time.Unix(100, 0), time.Unix(200, 0))
			require.NoError(t, err)
			require.NotNil(t, result["a"]["m"].USDPerMillion)
			assert.InDelta(t, 0.4, *result["a"]["m"].USDPerMillion, 1e-9)
			assert.Equal(t, int64(2), result["a"]["m"].Requests)
			assert.Equal(t, float64(1000), result["a"]["m"].Tokens)
			assert.InDelta(t, 2, *result["b"]["m"].USDPerMillion, 1e-9)
			assert.Zero(t, *result["a"]["free"].USDPerMillion)
			assert.NotContains(t, result["a"], "no-tokens")
		})
	}
}

func TestGroupCacheStatsDatabaseMatrix(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 is not configured")
				}
				driver = mysql.Open(dsn)
			case "postgres":
				dsn := os.Getenv("TEST_POSTGRES_DSN")
				if dsn == "" {
					t.Skip("TEST_POSTGRES_DSN is not configured")
				}
				driver = postgres.Open(dsn)
			}
			db, err := gorm.Open(driver, &gorm.Config{NamingStrategy: schema.NamingStrategy{TablePrefix: fmt.Sprintf("gs_%d_", time.Now().UnixNano())}})
			require.NoError(t, err)
			originalDB, originalLogDB := DB, LOG_DB
			DB, LOG_DB = db, db
			t.Cleanup(func() { DB, LOG_DB = originalDB, originalLogDB })
			channelTable := db.NamingStrategy.TableName("channels")
			logTable := db.NamingStrategy.TableName("logs")
			type sample struct {
				Id           int
				ChannelId    int
				PromptTokens int64
				Other        string
				Type         int
				TokenId      int
				CreatedAt    int64
			}
			require.NoError(t, db.Table(channelTable).AutoMigrate(&StatusChannel{}))
			require.NoError(t, db.Table(logTable).AutoMigrate(&sample{}))
			t.Cleanup(func() {
				require.NoError(t, db.Migrator().DropTable(channelTable, logTable))
				sqlDB, err := db.DB()
				require.NoError(t, err)
				require.NoError(t, sqlDB.Close())
			})
			require.NoError(t, db.Table(channelTable).Create(&[]StatusChannel{
				{Id: 1, Group: "renamed", Status: 1}, {Id: 2, Group: "automatic-down", Status: 3},
				{Id: 3, Group: "manual-off", Status: 2}, {Id: 4, Group: "unused", Status: 1},
			}).Error)
			require.NoError(t, db.Table(logTable).Create(&[]sample{
				{ChannelId: 1, PromptTokens: 100, Other: `{"cache_tokens":80}`, Type: 2, TokenId: 1, CreatedAt: 150},
				{ChannelId: 1, PromptTokens: 300, Other: `{"cache_tokens":0}`, Type: 2, TokenId: 1, CreatedAt: 150},
				{ChannelId: 1, PromptTokens: 20, Other: `{"cache_tokens":50,"usage_semantic":"anthropic","cache_creation_tokens":30}`, Type: 2, TokenId: 1, CreatedAt: 150},
				{ChannelId: 1, PromptTokens: 10, Other: `{"cache_tokens":40,"input_tokens_total":100,"claude":true}`, Type: 2, TokenId: 1, CreatedAt: 150},
				{ChannelId: 2, PromptTokens: 50, Other: `{}`, Type: 2, TokenId: 1, CreatedAt: 150},
				{ChannelId: 1, PromptTokens: 999, Other: `{}`, Type: 2, TokenId: 0, CreatedAt: 150},
				{ChannelId: 1, PromptTokens: 999, Other: `{}`, Type: 5, TokenId: 1, CreatedAt: 150},
				{ChannelId: 1, PromptTokens: 999, Other: `{}`, Type: 2, TokenId: 1, CreatedAt: 99},
				{ChannelId: 1, PromptTokens: 999, Other: `{}`, Type: 2, TokenId: 1, CreatedAt: 200},
				{ChannelId: 1, PromptTokens: 999, Other: `invalid`, Type: 2, TokenId: 1, CreatedAt: 150},
				{ChannelId: 1, PromptTokens: 0, Other: `{}`, Type: 2, TokenId: 1, CreatedAt: 150},
				{ChannelId: 3, PromptTokens: 999, Other: `{}`, Type: 2, TokenId: 1, CreatedAt: 150},
			}).Error)
			channels, err := GetStatusChannels(context.Background())
			require.NoError(t, err)
			require.Len(t, channels, 3)
			stats, err := GetGroupCacheStats(context.Background(), channels, time.Unix(100, 0), time.Unix(200, 0))
			require.NoError(t, err)
			assert.Equal(t, GroupCacheStats{Requests: 4, CachedRequests: 3, InputTokens: 600, CachedTokens: 170}, stats["renamed"])
			assert.Equal(t, GroupCacheStats{Requests: 1, InputTokens: 50}, stats["automatic-down"])
			assert.Equal(t, GroupCacheStats{}, stats["unused"])
			assert.NotContains(t, stats, "manual-off")
		})
	}
}