package helper

import (
	"net/http"
	"net/http/httptest"
	"testing"

	"github.com/QuantumNous/new-api/common"
	"github.com/QuantumNous/new-api/dto"
	"github.com/QuantumNous/new-api/model"
	relaycommon "github.com/QuantumNous/new-api/relay/common"
	"github.com/QuantumNous/new-api/setting/ratio_setting"
	"github.com/QuantumNous/new-api/types"
	"github.com/glebarez/sqlite"
	"github.com/gin-gonic/gin"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
	"gorm.io/gorm"
)

const testModelUsePriceSwitch = "test-useprice-switch-model"

func setupUsePriceSwitchTest(t *testing.T) {
	t.Helper()

	// 初始化 SQLite 内存数据库
	db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{})
	require.NoError(t, err)
	sqlDB, _ := db.DB()
	sqlDB.SetMaxOpenConns(1)

	origDB := model.DB
	model.DB = db
	common.UsingSQLite = true
	common.RedisEnabled = false

	require.NoError(t, db.AutoMigrate(&model.ChannelPricing{}))

	// 全局定价：按次计费（modelPrice 存在 → usePrice=true）
	require.NoError(t, ratio_setting.UpdateModelPriceByJSONString(`{"`+testModelUsePriceSwitch+`":0.5}`))
	// 全局 modelRatio 也设上（渠道按量计费时会用到）
	require.NoError(t, ratio_setting.UpdateModelRatioByJSONString(`{"`+testModelUsePriceSwitch+`":15}`))
	require.NoError(t, ratio_setting.UpdateCompletionRatioByJSONString(`{"`+testModelUsePriceSwitch+`":3}`))
	// 全局扩展比率
	require.NoError(t, ratio_setting.UpdateCacheRatioByJSONString(`{"`+testModelUsePriceSwitch+`":0.1}`))
	require.NoError(t, ratio_setting.UpdateCreateCacheRatioByJSONString(`{"`+testModelUsePriceSwitch+`":1.25}`))
	require.NoError(t, ratio_setting.UpdateImageRatioByJSONString(`{"`+testModelUsePriceSwitch+`":2.0}`))
	require.NoError(t, ratio_setting.UpdateAudioRatioByJSONString(`{"`+testModelUsePriceSwitch+`":5.0}`))
	require.NoError(t, ratio_setting.UpdateAudioCompletionRatioByJSONString(`{"`+testModelUsePriceSwitch+`":3.0}`))

	// 渠道定价：按量计费（QuotaTypeByTokens），扩展比率全部为 0（用全局）
	cp := &model.ChannelPricing{
		ModelName:       testModelUsePriceSwitch,
		ChannelId:       9901,
		QuotaType:       model.QuotaTypeByTokens,
		ModelRatio:      15,
		CompletionRatio: 3,
		// 扩展比率全部为 0 → 意味着回退全局
	}
	require.NoError(t, cp.Insert())

	t.Cleanup(func() {
		model.DB = origDB
		sqlDB.Close()
		// 清理全局定价
		ratio_setting.UpdateModelPriceByJSONString(`{}`)
		ratio_setting.UpdateModelRatioByJSONString(`{}`)
		ratio_setting.UpdateCompletionRatioByJSONString(`{}`)
		ratio_setting.UpdateCacheRatioByJSONString(`{}`)
		ratio_setting.UpdateCreateCacheRatioByJSONString(`{}`)
		ratio_setting.UpdateImageRatioByJSONString(`{}`)
		ratio_setting.UpdateAudioRatioByJSONString(`{}`)
		ratio_setting.UpdateAudioCompletionRatioByJSONString(`{}`)
	})
}

func buildTestContext(t *testing.T) *gin.Context {
	t.Helper()
	gin.SetMode(gin.TestMode)
	w := httptest.NewRecorder()
	c, _ := gin.CreateTestContext(w)
	c.Request, _ = http.NewRequest("POST", "/v1/chat/completions", nil)
	return c
}

// TestGlobalRatiosFallbackWhenGlobalUsePriceTrue 验证核心场景：
// 全局按次计费 → 渠道按量计费 → PriceData 中扩展比率仍应被全局默认值填充
// 而不是保持 0（否则后续 UpdatePriceDataForChannelPricing 切换为按量时会丢失比率）
func TestGlobalRatiosFallbackWhenGlobalUsePriceTrue(t *testing.T) {
	setupUsePriceSwitchTest(t)
	c := buildTestContext(t)

	// 模拟渠道选择后的场景：ChannelMeta 不为空，能查到渠道定价
	info := &relaycommon.RelayInfo{
		OriginModelName: testModelUsePriceSwitch,
		ChannelMeta: &relaycommon.ChannelMeta{
			ChannelId: 9901,
		},
		UserSetting: dto.UserSetting{},
	}
	meta := &types.TokenCountMeta{}

	priceData, err := ModelPriceHelper(c, info, 100, meta)
	require.NoError(t, err)

	// 断言：渠道按量计费覆盖了全局按次，扩展比率必须被全局默认值填充，不能是 0
	assert.Equal(t, 0.1, priceData.CacheRatio, "CacheRatio 应从全局回退填充")
	assert.Equal(t, 1.25, priceData.CacheCreationRatio, "CacheCreationRatio 应从全局回退填充")
	assert.Equal(t, 2.0, priceData.ImageRatio, "ImageRatio 应从全局回退填充")
	assert.Equal(t, 5.0, priceData.AudioRatio, "AudioRatio 应从全局回退填充")
	assert.Equal(t, 3.0, priceData.AudioCompletionRatio, "AudioCompletionRatio 应从全局回退填充")
}

// TestGlobalRatiosFallbackWithoutChannelPricing 验证无渠道定价时全局比率也正常回退
func TestGlobalRatiosFallbackWithoutChannelPricing(t *testing.T) {
	setupUsePriceSwitchTest(t)
	c := buildTestContext(t)

	// 无 ChannelMeta → 渠道定价不可用，走全局
	info := &relaycommon.RelayInfo{
		OriginModelName: testModelUsePriceSwitch,
		UserSetting:     dto.UserSetting{},
	}
	meta := &types.TokenCountMeta{}

	priceData, err := ModelPriceHelper(c, info, 100, meta)
	require.NoError(t, err)

	// 全局是按次计费，UsePrice=true，但比率仍应被填充
	assert.True(t, priceData.UsePrice)
	assert.Equal(t, 0.1, priceData.CacheRatio, "CacheRatio 应从全局回退填充")
	assert.Equal(t, 2.0, priceData.ImageRatio, "ImageRatio 应从全局回退填充")
}

// TestChannelRatiosOverrideGlobal 验证渠道非零值优先于全局
func TestChannelRatiosOverrideGlobal(t *testing.T) {
	setupUsePriceSwitchTest(t)

	// 渠道定价设了自定义 cacheRatio=0.5
	cp := &model.ChannelPricing{
		ModelName:       testModelUsePriceSwitch,
		ChannelId:       9902,
		QuotaType:       model.QuotaTypeByTokens,
		ModelRatio:      15,
		CompletionRatio: 3,
		CacheRatio:      0.5, // 渠道自定义，覆盖全局 0.1
		// 其他扩展比率保持 0，用全局
	}
	require.NoError(t, cp.Insert())

	c := buildTestContext(t)
	info := &relaycommon.RelayInfo{
		OriginModelName: testModelUsePriceSwitch,
		ChannelMeta: &relaycommon.ChannelMeta{
			ChannelId: 9902,
		},
		UserSetting: dto.UserSetting{},
	}

	priceData, err := ModelPriceHelper(c, info, 100, &types.TokenCountMeta{})
	require.NoError(t, err)

	assert.Equal(t, 0.5, priceData.CacheRatio, "渠道自定义 CacheRatio 应覆盖全局")
	assert.Equal(t, 2.0, priceData.ImageRatio, "ImageRatio 未设渠道值，应从全局回退")
	assert.Equal(t, 5.0, priceData.AudioRatio, "AudioRatio 未设渠道值，应从全局回退")
}