package controller

import (
	"maps"

	"github.com/QuantumNous/new-api/common"
	"github.com/QuantumNous/new-api/model"
	"github.com/QuantumNous/new-api/service"
	"github.com/QuantumNous/new-api/setting/billing_setting"
	"github.com/QuantumNous/new-api/setting/ratio_setting"

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

func filterPricingByUsableGroups(pricing []model.Pricing, usableGroup map[string]string) []model.Pricing {
	if len(pricing) == 0 {
		return pricing
	}
	if len(usableGroup) == 0 {
		return []model.Pricing{}
	}

	filtered := make([]model.Pricing, 0, len(pricing))
	for _, item := range pricing {
		if common.StringsContains(item.EnableGroup, "all") {
			filtered = append(filtered, item)
			continue
		}
		for _, group := range item.EnableGroup {
			if _, ok := usableGroup[group]; ok {
				filtered = append(filtered, item)
				break
			}
		}
	}
	return filtered
}

func GetPricing(c *gin.Context) {
	pricing := model.GetPricing()
	enabledGroups, err := model.GetEnabledGroupSet()
	if err != nil {
		common.ApiError(c, err)
		return
	}
	userId, exists := c.Get("id")
	usableGroup := map[string]string{}
	groupRatio := map[string]float64{}
	maps.Copy(groupRatio, ratio_setting.GetGroupRatioCopy())
	var group string
	if exists {
		user, err := model.GetUserCache(userId.(int))
		if err == nil {
			group = user.Group
			if err := model.LoadPersonalGroupRatios(c, user.Id); err != nil {
				common.ApiError(c, err)
				return
			}
			for g := range groupRatio {
				groupRatio[g] = model.EffectiveGroupRatio(c, group, g).GroupRatio
			}
		}
	}

	usableGroup = service.GetUserUsableGroups(group)
	for group := range usableGroup {
		if group == "auto" {
			continue
		}
		if _, ok := enabledGroups[group]; !ok {
			delete(usableGroup, group)
		}
	}
	if len(enabledGroups) == 0 {
		delete(usableGroup, "auto")
	}
	pricing = filterPricingByUsableGroups(pricing, usableGroup)
	pricing = append([]model.Pricing(nil), pricing...)
	for i := range pricing {
		for _, group := range pricing[i].EnableGroup {
			if _, visible := usableGroup[group]; !visible {
				continue
			}
			if expr, ok := billing_setting.GetGroupBillingExpr(group, pricing[i].ModelName); ok {
				if pricing[i].GroupBillingExpr == nil {
					pricing[i].GroupBillingExpr = make(map[string]string)
				}
				pricing[i].GroupBillingExpr[group] = expr
			}
		}
	}
	// check groupRatio contains usableGroup
	for group := range ratio_setting.GetGroupRatioCopy() {
		_, usable := usableGroup[group]
		_, enabled := enabledGroups[group]
		if !usable || !enabled {
			delete(groupRatio, group)
		}
	}

	c.JSON(200, gin.H{
		"success":            true,
		"data":               pricing,
		"vendors":            model.GetVendors(),
		"group_ratio":        groupRatio,
		"usable_group":       usableGroup,
		"supported_endpoint": model.GetSupportedEndpointMap(),
		"auto_groups":        service.GetUserAutoGroup(group),
		"pricing_version":    "a42d372ccf0b5dd13ecf71203521f9d2",
	})
}

func ResetModelRatio(c *gin.Context) {
	defaultStr := ratio_setting.DefaultModelRatio2JSONString()
	err := model.UpdateOption("ModelRatio", defaultStr)
	if err != nil {
		common.ApiError(c, err)
		return
	}
	err = ratio_setting.UpdateModelRatioByJSONString(defaultStr)
	if err != nil {
		common.ApiError(c, err)
		return
	}
	common.ApiSuccessT(c, "Model ratios reset", nil)
}