package model

import (
	"testing"

	"github.com/QuantumNous/new-api/setting/ratio_setting"
	"github.com/glebarez/sqlite"
	"github.com/stretchr/testify/require"
	"gorm.io/gorm"
)

const testPricingModel = "test-pricing-model-apply"

func setupPricingTest(t *testing.T) {
	t.Helper()
	db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{})
	require.NoError(t, err)
	sqlDB, _ := db.DB()
	sqlDB.SetMaxOpenConns(1)

	origDB := DB
	DB = db
	require.NoError(t, db.AutoMigrate(&ChannelPricing{}))

	// 全局默认定价
	require.NoError(t, ratio_setting.UpdateModelRatioByJSONString(`{"`+testPricingModel+`":10}`))
	require.NoError(t, ratio_setting.UpdateCompletionRatioByJSONString(`{"`+testPricingModel+`":3}`))
	require.NoError(t, ratio_setting.UpdateCacheRatioByJSONString(`{"`+testPricingModel+`":0.5}`))
	require.NoError(t, ratio_setting.UpdateCreateCacheRatioByJSONString(`{"`+testPricingModel+`":0.75}`))

	t.Cleanup(func() {
		DB = origDB
		sqlDB.Close()
		ratio_setting.UpdateModelRatioByJSONString(`{}`)
		ratio_setting.UpdateCompletionRatioByJSONString(`{}`)
		ratio_setting.UpdateCacheRatioByJSONString(`{}`)
		ratio_setting.UpdateCreateCacheRatioByJSONString(`{}`)
		ratio_setting.UpdateModelPriceByJSONString(`{}`)
	})
}

func TestApplyBestChannelPricing_NoChannelPricing(t *testing.T) {
	setupPricingTest(t)

	p := &Pricing{}
	applyBestChannelPricing(p, nil, testPricingModel)

	require.Equal(t, 0, p.QuotaType, "应使用全局按量计费")
	require.Equal(t, 10.0, p.ModelRatio, "应使用全局 model_ratio")
	require.Equal(t, 3.0, p.CompletionRatio, "应使用全局 completion_ratio")
	require.Equal(t, 0.5, p.CacheRatio, "应使用全局 cache_ratio")
	require.Equal(t, 0.75, p.CacheCreationRatio, "应使用全局 cache_creation_ratio")
}

func TestApplyBestChannelPricing_PerTokenCheapest(t *testing.T) {
	setupPricingTest(t)

	// 渠道A: model_ratio=8（更便宜）
	require.NoError(t, (&ChannelPricing{
		ModelName: testPricingModel, ChannelId: 1,
		QuotaType: QuotaTypeByTokens, ModelRatio: 8, CompletionRatio: 2,
		CacheRatio: 0.3, CacheCreationRatio: 0.6,
	}).Insert())
	// 渠道B: model_ratio=12（更贵）
	require.NoError(t, (&ChannelPricing{
		ModelName: testPricingModel, ChannelId: 2,
		QuotaType: QuotaTypeByTokens, ModelRatio: 12, CompletionRatio: 4,
		CacheRatio: 0.8, CacheCreationRatio: 1.0,
	}).Insert())

	cps := []ChannelPricing{
		{ModelRatio: 8, CompletionRatio: 2, CacheRatio: 0.3, CacheCreationRatio: 0.6},
		{ModelRatio: 12, CompletionRatio: 4, CacheRatio: 0.8, CacheCreationRatio: 1.0},
	}

	p := &Pricing{}
	applyBestChannelPricing(p, cps, testPricingModel)

	require.Equal(t, 0, p.QuotaType)
	require.Equal(t, 8.0, p.ModelRatio, "应选最便宜的渠道A")
	require.Equal(t, 2.0, p.CompletionRatio)
	require.Equal(t, 0.3, p.CacheRatio)
	require.Equal(t, 0.6, p.CacheCreationRatio)
}

func TestApplyBestChannelPricing_ExtendedRatioZeroFallback(t *testing.T) {
	setupPricingTest(t)

	// 渠道定价中扩展比率为 0，应回退到全局值
	require.NoError(t, (&ChannelPricing{
		ModelName: testPricingModel, ChannelId: 1,
		QuotaType: QuotaTypeByTokens, ModelRatio: 5, CompletionRatio: 1,
		CacheRatio: 0, CacheCreationRatio: 0, // 0 = 未设置
	}).Insert())

	cps := []ChannelPricing{
		{ModelRatio: 5, CompletionRatio: 1, CacheRatio: 0, CacheCreationRatio: 0},
	}

	p := &Pricing{}
	applyBestChannelPricing(p, cps, testPricingModel)

	require.Equal(t, 0.5, p.CacheRatio, "cache_ratio=0 应回退全局 0.5")
	require.Equal(t, 0.75, p.CacheCreationRatio, "cache_creation_ratio=0 应回退全局 0.75")
}

func TestApplyBestChannelPricing_PerCallCheapest(t *testing.T) {
	setupPricingTest(t)
	require.NoError(t, ratio_setting.UpdateModelPriceByJSONString(`{"`+testPricingModel+`":0.8}`))

	// 只有按次计费的渠道
	cps := []ChannelPricing{
		{QuotaType: QuotaTypeByCall, ModelPrice: 0.3},
		{QuotaType: QuotaTypeByCall, ModelPrice: 0.5},
	}

	p := &Pricing{}
	applyBestChannelPricing(p, cps, testPricingModel)

	require.Equal(t, 1, p.QuotaType)
	require.Equal(t, 0.3, p.ModelPrice, "应选最便宜的按次渠道")
}
