package model

import (
	"testing"

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

func setupChannelPricingDB(t *testing.T) *gorm.DB {
	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
	common.UsingSQLite = true
	common.RedisEnabled = false

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

	t.Cleanup(func() {
		DB = origDB
		sqlDB.Close()
	})
	return db
}

func TestCacheWriteThrough_Insert(t *testing.T) {
	setupChannelPricingDB(t)

	cp := &ChannelPricing{
		ModelName:  "test-insert-model",
		ChannelId:  9001,
		QuotaType:  QuotaTypeByTokens,
		ModelRatio: 1.0,
		CacheRatio: 0.8,
		ImageRatio: 1.2,
	}
	require.NoError(t, cp.Insert())
	defer cp.Delete()

	found, ok := GetEffectivePricing("test-insert-model", 9001)
	assert.True(t, ok)
	assert.Equal(t, 0.8, found.CacheRatio)
	assert.Equal(t, 1.2, found.ImageRatio)
}

func TestCacheWriteThrough_Update(t *testing.T) {
	setupChannelPricingDB(t)

	cp := &ChannelPricing{
		ModelName:  "test-update-model",
		ChannelId:  9002,
		QuotaType:  QuotaTypeByTokens,
		ModelRatio: 1.0,
	}
	require.NoError(t, cp.Insert())
	defer cp.Delete()

	cp.CacheRatio = 0.9
	cp.AudioRatio = 1.5
	require.NoError(t, cp.Update())

	found, ok := GetEffectivePricing("test-update-model", 9002)
	assert.True(t, ok)
	assert.Equal(t, 0.9, found.CacheRatio)
	assert.Equal(t, 1.5, found.AudioRatio)
}

func TestCacheWriteThrough_Delete(t *testing.T) {
	setupChannelPricingDB(t)

	cp := &ChannelPricing{
		ModelName:  "test-delete-model",
		ChannelId:  9003,
		QuotaType:  QuotaTypeByTokens,
		ModelRatio: 1.0,
	}
	require.NoError(t, cp.Insert())

	require.NoError(t, cp.Delete())

	found, ok := GetEffectivePricing("test-delete-model", 9003)
	assert.False(t, ok)
	assert.Nil(t, found)
}

func TestExtendedFields_DefaultZero(t *testing.T) {
	setupChannelPricingDB(t)

	cp := &ChannelPricing{
		ModelName:  "test-default-model",
		ChannelId:  9004,
		QuotaType:  QuotaTypeByTokens,
		ModelRatio: 1.0,
	}
	require.NoError(t, cp.Insert())
	defer cp.Delete()

	found, ok := GetEffectivePricing("test-default-model", 9004)
	assert.True(t, ok)
	assert.Equal(t, 0.0, found.CacheRatio)
	assert.Equal(t, 0.0, found.CacheCreationRatio)
	assert.Equal(t, 0.0, found.ImageRatio)
	assert.Equal(t, 0.0, found.AudioRatio)
	assert.Equal(t, 0.0, found.AudioCompletionRatio)
}

func TestGetEffectivePricing_NotFound(t *testing.T) {
	setupChannelPricingDB(t)

	found, ok := GetEffectivePricing("nonexistent-model-xyz", 99999)
	assert.False(t, ok)
	assert.Nil(t, found)
}

// === 默认通道测试 ===

func TestDefaultChannel_SetAndGet(t *testing.T) {
	setupChannelPricingDB(t)

	// 创建两条定价记录
	cp1 := &ChannelPricing{ModelName: "default-test-model", ChannelId: 100, QuotaType: QuotaTypeByTokens, ModelRatio: 1.0}
	cp2 := &ChannelPricing{ModelName: "default-test-model", ChannelId: 200, QuotaType: QuotaTypeByTokens, ModelRatio: 2.0}
	require.NoError(t, cp1.Insert())
	require.NoError(t, cp2.Insert())
	t.Cleanup(func() { cp1.Delete(); cp2.Delete() })

	// 初始没有默认
	_, ok := GetDefaultChannelId("default-test-model")
	assert.False(t, ok)

	// 设置通道 100 为默认
	require.NoError(t, SetDefaultChannel("default-test-model", 100))
	id, ok := GetDefaultChannelId("default-test-model")
	assert.True(t, ok)
	assert.Equal(t, 100, id)

	// 切换默认到通道 200
	require.NoError(t, SetDefaultChannel("default-test-model", 200))
	id, ok = GetDefaultChannelId("default-test-model")
	assert.True(t, ok)
	assert.Equal(t, 200, id)

	// 验证旧的默认标记被清除
	found, _ := GetEffectivePricing("default-test-model", 100)
	assert.False(t, found.IsDefault)
	found, _ = GetEffectivePricing("default-test-model", 200)
	assert.True(t, found.IsDefault)
}

func TestDefaultChannel_Clear(t *testing.T) {
	setupChannelPricingDB(t)

	cp := &ChannelPricing{ModelName: "clear-test-model", ChannelId: 300, QuotaType: QuotaTypeByTokens, ModelRatio: 1.0}
	require.NoError(t, cp.Insert())
	t.Cleanup(func() { cp.Delete() })

	require.NoError(t, SetDefaultChannel("clear-test-model", 300))
	_, ok := GetDefaultChannelId("clear-test-model")
	assert.True(t, ok)

	require.NoError(t, ClearDefaultChannel("clear-test-model"))
	_, ok = GetDefaultChannelId("clear-test-model")
	assert.False(t, ok)
}

func TestDefaultChannel_DeleteClearsCache(t *testing.T) {
	setupChannelPricingDB(t)

	cp := &ChannelPricing{ModelName: "delete-test-model", ChannelId: 400, QuotaType: QuotaTypeByTokens, ModelRatio: 1.0, IsDefault: true}
	require.NoError(t, cp.Insert())

	id, ok := GetDefaultChannelId("delete-test-model")
	assert.True(t, ok)
	assert.Equal(t, 400, id)

	// 删除记录后缓存应清除
	require.NoError(t, cp.Delete())
	_, ok = GetDefaultChannelId("delete-test-model")
	assert.False(t, ok)
}

func TestDefaultChannel_LoadCache(t *testing.T) {
	db := setupChannelPricingDB(t)

	// 直接插入数据（绕过缓存）
	db.Create(&ChannelPricing{ModelName: "load-model", ChannelId: 500, ModelRatio: 1.0, IsDefault: true})
	db.Create(&ChannelPricing{ModelName: "load-model", ChannelId: 501, ModelRatio: 2.0, IsDefault: false})
	db.Create(&ChannelPricing{ModelName: "other-model", ChannelId: 502, ModelRatio: 1.0, IsDefault: true})

	// 全量加载缓存
	LoadChannelPricingCache()

	id, ok := GetDefaultChannelId("load-model")
	assert.True(t, ok)
	assert.Equal(t, 500, id)

	id, ok = GetDefaultChannelId("other-model")
	assert.True(t, ok)
	assert.Equal(t, 502, id)

	// 无默认的模型
	_, ok = GetDefaultChannelId("nonexistent")
	assert.False(t, ok)
}

func TestDefaultChannel_InsertWithDefault(t *testing.T) {
	setupChannelPricingDB(t)

	cp := &ChannelPricing{ModelName: "insert-default-model", ChannelId: 600, ModelRatio: 1.0, IsDefault: true}
	require.NoError(t, cp.Insert())
	t.Cleanup(func() { cp.Delete() })

	id, ok := GetDefaultChannelId("insert-default-model")
	assert.True(t, ok)
	assert.Equal(t, 600, id)
}

func TestDefaultChannel_UpdateWithDefault(t *testing.T) {
	setupChannelPricingDB(t)

	cp1 := &ChannelPricing{ModelName: "update-default-model", ChannelId: 700, ModelRatio: 1.0, IsDefault: true}
	cp2 := &ChannelPricing{ModelName: "update-default-model", ChannelId: 701, ModelRatio: 2.0}
	require.NoError(t, cp1.Insert())
	require.NoError(t, cp2.Insert())
	t.Cleanup(func() { cp1.Delete(); cp2.Delete() })

	// cp1 是默认，通过 Update 把 cp2 设为默认
	cp2.IsDefault = true
	require.NoError(t, cp2.Update())

	id, ok := GetDefaultChannelId("update-default-model")
	assert.True(t, ok)
	assert.Equal(t, 701, id)
}
