package model

import (
	"context"
	"fmt"
	"os"
	"testing"
	"time"

	"github.com/QuantumNous/new-api/common"
	"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 TestGroupModelPerformanceDatabaseMatrix(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("performance_%d_", time.Now().UnixNano())}})
			require.NoError(t, err)
			conn, err := db.DB()
			require.NoError(t, err)
			conn.SetMaxOpenConns(1)
			previousDB, previousLogs := DB, LOG_DB
			DB, LOG_DB = db, db
			t.Cleanup(func() {
				for _, table := range []any{&Log{}, &Ability{}, &Channel{}} {
					require.NoError(t, db.Migrator().DropTable(table))
				}
				DB, LOG_DB = previousDB, previousLogs
				require.NoError(t, conn.Close())
			})
			query := "select version()"
			if dialect == "sqlite" {
				query = "select sqlite_version()"
			}
			var version string
			require.NoError(t, db.Raw(query).Scan(&version).Error)
			t.Logf("database version: %s", version)
			require.NoError(t, db.AutoMigrate(&Log{}, &Ability{}, &Channel{}))
			require.NoError(t, db.Create(&Channel{Id: 1, Status: common.ChannelStatusEnabled, Group: "primary,backup"}).Error)
			require.NoError(t, db.Create(&Channel{Id: 2, Status: common.ChannelStatusEnabled, Group: "primary"}).Error)
			require.NoError(t, db.Create(&Channel{Id: 3, Status: common.ChannelStatusAutoDisabled, Group: "primary"}).Error)
			for _, channel := range []int{1, 2, 3} {
				require.NoError(t, db.Create(&Ability{ChannelId: channel, Group: "primary", Model: "m", Enabled: true}).Error)
			}
			ctx := context.Background()
			empty, err := GetGroupModelPerformance(ctx, "primary", "m")
			require.NoError(t, err)
			assert.Nil(t, empty.Throughput)
			assert.Nil(t, empty.FirstToken)
			logs := []Log{
				{ChannelId: 1, Type: LogTypeConsume, Content: "模型测试", ModelName: "m", IsStream: true, CreatedAt: 1, UseTime: 10, CompletionTokens: 100, Other: `{"frt":1000}`},
				{ChannelId: 1, Type: LogTypeConsume, Content: "Model test", ModelName: "m", IsStream: true, CreatedAt: 2, UseTime: 1, CompletionTokens: 20, Other: `{"frt":500,"latency_ms":500}`},
				{ChannelId: 2, Type: LogTypeConsume, Content: "Model test", ModelName: "m", IsStream: true, CreatedAt: 2, UseTime: 10, CompletionTokens: 200, Other: `{"frt":1500}`},
				{ChannelId: 3, Type: LogTypeConsume, Content: "Model test", ModelName: "m", IsStream: true, CreatedAt: 3, UseTime: 1, CompletionTokens: 999, Other: `{"frt":1}`},
				{ChannelId: 1, Type: LogTypeConsume, Content: "Model test", ModelName: "other", IsStream: true, CreatedAt: 3, UseTime: 1, CompletionTokens: 999, Other: `{"frt":1}`},
				{ChannelId: 1, Type: LogTypeConsume, Content: "Model test", ModelName: "m", IsStream: true, CreatedAt: 4, UseTime: 1, CompletionTokens: 999, Other: `{"frt":1,"stream_status":{"status":"error"}}`},
			}
			require.NoError(t, db.Create(&logs).Error)
			fallback, err := GetGroupModelPerformance(ctx, "primary", "m")
			require.NoError(t, err)
			require.NotNil(t, fallback.Throughput)
			assert.Equal(t, ModelPerformanceMetric{Value: 30, Samples: 2, Source: "test", UpdatedAt: 2}, *fallback.Throughput)
			assert.Equal(t, 1000.0, fallback.FirstToken.Value)
			calls := make([]Log, 0, 31)
			for i := range 31 {
				calls = append(calls, Log{ChannelId: 1, Type: LogTypeConsume, TokenId: 8, Group: "primary", ModelName: "m", IsStream: true, CreatedAt: int64(100 + i), UseTime: i + 1, CompletionTokens: (i + 1) * (i + 1), Other: fmt.Sprintf(`{"frt":%d}`, (i+1)*100)})
			}
			require.NoError(t, db.Create(&calls).Error)
			// Invalid, another group's, and synthetic requests must not displace real samples.
			noise := []Log{
				{Type: LogTypeConsume, TokenId: 8, Group: "backup", ModelName: "m", IsStream: true, CreatedAt: 999, UseTime: 1, CompletionTokens: 999, Other: `{"frt":1}`},
				{Type: LogTypeConsume, TokenId: 8, Group: "primary", ModelName: "other", IsStream: true, CreatedAt: 999, UseTime: 1, CompletionTokens: 999, Other: `{"frt":1}`},
				{Type: LogTypeConsume, TokenId: 8, Group: "primary", ModelName: "m", IsStream: true, CreatedAt: 999, UseTime: 1, CompletionTokens: 999, Other: `{"frt":1,"stream_status":{"status":"error"}}`},
				{Type: LogTypeConsume, TokenId: 8, Group: "primary", ModelName: "m", IsStream: true, CreatedAt: 999, UseTime: 1, CompletionTokens: 999, Other: `{"frt":1,"stream_status":{"status":"ok","end_reason":"client_gone"}}`},
			}
			require.NoError(t, db.Create(&noise).Error)
			actual, err := GetGroupModelPerformance(ctx, "primary", "m")
			require.NoError(t, err)
			assert.Equal(t, ModelPerformanceMetric{Value: 16.5, Samples: 30, Source: "calls", UpdatedAt: 130}, *actual.Throughput)
			assert.Equal(t, ModelPerformanceMetric{Value: 1650, Samples: 30, Source: "calls", UpdatedAt: 130}, *actual.FirstToken)
			// Synchronous requests contribute throughput, not fabricated first-token timing.
			require.NoError(t, db.Create(&Log{Type: LogTypeConsume, TokenId: 8, Group: "primary", ModelName: "m", CreatedAt: 1000, CompletionTokens: 10, Other: `{"latency_ms":500,"frt":500}`}).Error)
			actual, err = GetGroupModelPerformance(ctx, "primary", "m")
			require.NoError(t, err)
			assert.InDelta(t, (16.5*30-2+20)/30, actual.Throughput.Value, 0.0001)
			assert.Equal(t, 1650.0, actual.FirstToken.Value)
			require.NoError(t, db.Model(&Log{}).Where(map[string]any{"group": "primary", "model_name": "m"}).Where("token_id > 0").Update("is_stream", false).Error)
			partial, err := GetGroupModelPerformance(ctx, "primary", "m")
			require.NoError(t, err)
			assert.Equal(t, "calls", partial.Throughput.Source)
			assert.Equal(t, "test", partial.FirstToken.Source)
			assert.Equal(t, 1000.0, partial.FirstToken.Value)
		})
	}
}