package helper

import (
	"fmt"

	"github.com/QuantumNous/new-api/common"
	"github.com/QuantumNous/new-api/constant"
	"github.com/QuantumNous/new-api/logger"
	"github.com/QuantumNous/new-api/model"
	relaycommon "github.com/QuantumNous/new-api/relay/common"
	"github.com/QuantumNous/new-api/setting/operation_setting"
	"github.com/QuantumNous/new-api/setting/ratio_setting"
	"github.com/QuantumNous/new-api/types"

	"github.com/gin-gonic/gin"
)

// HandleGroupRatio checks for "auto_group" in the context and updates the group ratio and relayInfo.UsingGroup if present
func HandleGroupRatio(ctx *gin.Context, relayInfo *relaycommon.RelayInfo) types.GroupRatioInfo {
	groupRatioInfo := types.GroupRatioInfo{
		GroupRatio:        1.0, // default ratio
		GroupSpecialRatio: -1,
	}

	// check auto group
	autoGroup, exists := ctx.Get("auto_group")
	if exists {
		logger.LogDebug(ctx, fmt.Sprintf("final group: %s", autoGroup))
		relayInfo.UsingGroup = autoGroup.(string)
	}

	// check user group special ratio
	userGroupRatio, ok := ratio_setting.GetGroupGroupRatio(relayInfo.UserGroup, relayInfo.UsingGroup)
	if ok {
		// user group special ratio
		groupRatioInfo.GroupSpecialRatio = userGroupRatio
		groupRatioInfo.GroupRatio = userGroupRatio
		groupRatioInfo.HasSpecialRatio = true
	} else {
		// normal group ratio
		groupRatioInfo.GroupRatio = ratio_setting.GetGroupRatio(relayInfo.UsingGroup)
	}

	return groupRatioInfo
}

func ModelPriceHelper(c *gin.Context, info *relaycommon.RelayInfo, promptTokens int, meta *types.TokenCountMeta) (types.PriceData, error) {
	var modelPrice float64
	var usePrice bool
	var modelRatio float64
	var completionRatio float64
	var channelPricingFound bool
	var cacheRatio, cacheCreationRatio, imageRatio, audioRatio, audioCompletionRatio float64

	// 尝试获取渠道定价（优先于全局定价）
	// ChannelMeta 在 InitChannelMeta 之前为 nil，但 Distribute 已将 channelId 写入 context
	channelId := 0
	if info != nil && info.ChannelMeta != nil && info.ChannelId > 0 {
		channelId = info.ChannelId
	} else {
		channelId = common.GetContextKeyInt(c, constant.ContextKeyChannelId)
	}
	if channelId > 0 {
		cp, found := model.GetEffectivePricing(info.OriginModelName, channelId)
		if found && cp != nil {
			modelRatio = cp.ModelRatio
			completionRatio = cp.CompletionRatio
			modelPrice = cp.ModelPrice
			usePrice = cp.QuotaType == model.QuotaTypeByCall
			channelPricingFound = true
			// 渠道定价的扩展比率（非零值直接使用）
			cacheRatio = cp.CacheRatio
			cacheCreationRatio = cp.CacheCreationRatio
			imageRatio = cp.ImageRatio
			audioRatio = cp.AudioRatio
			audioCompletionRatio = cp.AudioCompletionRatio
		}
	}

	// 如果没有渠道定价，使用全局定价
	if !channelPricingFound {
		modelPrice, usePrice = ratio_setting.GetModelPrice(info.OriginModelName, false)
	}

	groupRatioInfo := HandleGroupRatio(c, info)

	var preConsumedQuota int
	var freeModel bool
	if !usePrice {
		preConsumedTokens := common.Max(promptTokens, common.PreConsumedQuota)
		if meta.MaxTokens != 0 {
			preConsumedTokens += meta.MaxTokens
		}
		// 如果没有找到渠道定价，获取全局 modelRatio 和 completionRatio
		if !channelPricingFound {
			var success bool
			var matchName string
			modelRatio, success, matchName = ratio_setting.GetModelRatio(info.OriginModelName)
			if !success {
				acceptUnsetRatio := false
				if info.UserSetting.AcceptUnsetRatioModel {
					acceptUnsetRatio = true
				}
				if !acceptUnsetRatio {
					return types.PriceData{}, fmt.Errorf("模型 %s 倍率或价格未配置，请联系管理员设置或开启自用模式；Model %s ratio or price not set, please set or start self-use mode", matchName, matchName)
				}
			}
			completionRatio = ratio_setting.GetCompletionRatio(info.OriginModelName)
		}
		ratio := modelRatio * groupRatioInfo.GroupRatio
		preConsumedQuota = int(float64(preConsumedTokens) * ratio)
	} else {
		if meta.ImagePriceRatio != 0 {
			modelPrice = modelPrice * meta.ImagePriceRatio
		}
		preConsumedQuota = int(modelPrice * common.QuotaPerUnit * groupRatioInfo.GroupRatio)
	}

	// 全局比率作为回退（仅当渠道未设置、值为 0 时生效）
	// 必须放在 usePrice 判断之外，因为 UpdatePriceDataForChannelPricing 可能改变 UsePrice
	if cacheRatio == 0 {
		cacheRatio, _ = ratio_setting.GetCacheRatio(info.OriginModelName)
	}
	if cacheCreationRatio == 0 {
		cacheCreationRatio, _ = ratio_setting.GetCreateCacheRatio(info.OriginModelName)
	}
	if imageRatio == 0 {
		imageRatio, _ = ratio_setting.GetImageRatio(info.OriginModelName)
	}
	if audioRatio == 0 {
		audioRatio = ratio_setting.GetAudioRatio(info.OriginModelName)
	}
	if audioCompletionRatio == 0 {
		audioCompletionRatio = ratio_setting.GetAudioCompletionRatio(info.OriginModelName)
	}

	// check if free model pre-consume is disabled
	if !operation_setting.GetQuotaSetting().EnableFreeModelPreConsume {
		// if model price or ratio is 0, do not pre-consume quota
		if groupRatioInfo.GroupRatio == 0 {
			preConsumedQuota = 0
			freeModel = true
		} else if usePrice {
			if modelPrice == 0 {
				preConsumedQuota = 0
				freeModel = true
			}
		} else {
			if modelRatio == 0 {
				preConsumedQuota = 0
				freeModel = true
			}
		}
	}

	priceData := types.PriceData{
		FreeModel:            freeModel,
		ModelPrice:           modelPrice,
		ModelRatio:           modelRatio,
		CompletionRatio:      completionRatio,
		GroupRatioInfo:       groupRatioInfo,
		UsePrice:             usePrice,
		QuotaToPreConsume:    preConsumedQuota,
		CacheRatio:           cacheRatio,
		ImageRatio:           imageRatio,
		AudioRatio:           audioRatio,
		AudioCompletionRatio: audioCompletionRatio,
		CacheCreationRatio:   cacheCreationRatio,
		CacheCreation5mRatio: cacheCreationRatio,
		CacheCreation1hRatio: cacheCreationRatio * types.ClaudeCacheCreation1hMultiplier,
	}

	if common.DebugEnabled {
		println(fmt.Sprintf("model_price_helper result: %s", priceData.ToSetting()))
	}
	info.PriceData = priceData
	return priceData, nil
}

// ModelPriceHelperPerCall 按次计费的 PriceHelper (MJ、Task)
func ModelPriceHelperPerCall(c *gin.Context, info *relaycommon.RelayInfo) types.PriceData {
	groupRatioInfo := HandleGroupRatio(c, info)

	modelPrice, success := ratio_setting.GetModelPrice(info.OriginModelName, true)
	// 如果没有配置价格，则使用默认价格
	if !success {
		defaultPrice, ok := ratio_setting.GetDefaultModelPriceMap()[info.OriginModelName]
		if !ok {
			modelPrice = 0.1
		} else {
			modelPrice = defaultPrice
		}
	}
	quota := int(modelPrice * common.QuotaPerUnit * groupRatioInfo.GroupRatio)

	// 免费模型检测（与 ModelPriceHelper 对齐）
	freeModel := false
	if !operation_setting.GetQuotaSetting().EnableFreeModelPreConsume {
		if groupRatioInfo.GroupRatio == 0 || modelPrice == 0 {
			quota = 0
			freeModel = true
		}
	}

	priceData := types.PriceData{
		FreeModel:      freeModel,
		ModelPrice:     modelPrice,
		Quota:          quota,
		GroupRatioInfo: groupRatioInfo,
	}
	return priceData
}

func ContainPriceOrRatio(modelName string) bool {
	_, ok := ratio_setting.GetModelPrice(modelName, false)
	if ok {
		return true
	}
	_, ok, _ = ratio_setting.GetModelRatio(modelName)
	if ok {
		return true
	}
	return false
}

// UpdatePriceDataForChannelPricing 在渠道选择后更新 PriceData 以使用渠道定价
// ModelPriceHelper 在渠道选择之前被调用，此时 ChannelMeta 为 nil，导致渠道定价无法使用。
// 需要在渠道选择后调用此函数更新价格数据，用于后续的日志记录和计费计算。
func UpdatePriceDataForChannelPricing(c *gin.Context, info *relaycommon.RelayInfo, channelId int) {
	if info == nil || channelId <= 0 {
		return
	}

	cp, found := model.GetEffectivePricing(info.OriginModelName, channelId)
	if !found {
		return
	}

	info.PriceData.ModelRatio = cp.ModelRatio
	info.PriceData.CompletionRatio = cp.CompletionRatio
	info.PriceData.UsePrice = cp.QuotaType == model.QuotaTypeByCall
	if info.PriceData.UsePrice {
		info.PriceData.ModelPrice = cp.ModelPrice
	} else {
		info.PriceData.ModelPrice = -1
	}

	info.PriceData.ApplyChannelPricingRatios(cp.CacheRatio, cp.CacheCreationRatio, cp.ImageRatio, cp.AudioRatio, cp.AudioCompletionRatio)

	if info.PriceData.UsePrice {
		info.PriceData.QuotaToPreConsume = int(
			cp.ModelPrice * common.QuotaPerUnit * info.PriceData.GroupRatioInfo.GroupRatio)
	} else {
		estimateTokens := info.GetEstimatePromptTokens()
		if estimateTokens > 0 {
			ratio := cp.ModelRatio * info.PriceData.GroupRatioInfo.GroupRatio
			info.PriceData.QuotaToPreConsume = int(float64(estimateTokens) * ratio)
		}
	}

	if common.DebugEnabled {
		println(fmt.Sprintf("[ChannelPricing] updatePriceData: model=%s channel=%d modelRatio=%.4f completionRatio=%.4f cacheRatio=%.4f imageRatio=%.4f audioRatio=%.4f",
			info.OriginModelName, channelId, cp.ModelRatio, cp.CompletionRatio, cp.CacheRatio, cp.ImageRatio, cp.AudioRatio))
	}
}
