package model

import (
	"context"
	"encoding/json"
	"fmt"
	"math"

	"github.com/QuantumNous/new-api/setting/ratio_setting"
	"github.com/QuantumNous/new-api/types"
	"github.com/gin-gonic/gin"
	"gorm.io/gorm/clause"
)

// Stored independently from User and UserSetting to prevent profile mass assignment.
type PersonalGroupRatio struct {
	UserID int    `gorm:"primaryKey;autoIncrement:false"`
	Ratios string `gorm:"type:text;not null"`
}

const personalRatioContextKey = "personal_group_ratios"

func GetPersonalGroupRatios(ctx context.Context, userID int) (map[string]float64, error) {
	var rows []PersonalGroupRatio
	if err := DB.WithContext(ctx).Where("user_id = ?", userID).Find(&rows).Error; err != nil {
		return nil, err
	}
	result := map[string]float64{}
	if len(rows) != 0 {
		if err := json.Unmarshal([]byte(rows[0].Ratios), &result); err != nil {
			return nil, err
		}
	}
	return result, nil
}

func SavePersonalGroupRatios(ctx context.Context, userID int, ratios map[string]float64) error {
	for group, ratio := range ratios {
		if !ratio_setting.ContainsGroupRatio(group) || group == "auto" || ratio < 0 || math.IsNaN(ratio) || math.IsInf(ratio, 0) {
			return fmt.Errorf("invalid group or ratio: %s", group)
		}
	}
	if ratios == nil {
		ratios = map[string]float64{}
	}
	body, err := json.Marshal(ratios)
	if err != nil {
		return err
	}
	row := PersonalGroupRatio{UserID: userID, Ratios: string(body)}
	if err = DB.WithContext(ctx).Clauses(clause.OnConflict{Columns: []clause.Column{{Name: "user_id"}}, DoUpdates: clause.AssignmentColumns([]string{"ratios"})}).Create(&row).Error; err != nil {
		return err
	}
	// Pricing is loaded from the database at request admission, never from a stale
	// process/Redis cache. The request snapshot survives retries and price updates.
	return invalidateUserCache(userID)
}

func LoadPersonalGroupRatios(c *gin.Context, userID int) error {
	ratios, err := GetPersonalGroupRatios(c.Request.Context(), userID)
	if err == nil {
		c.Set(personalRatioContextKey, ratios)
	}
	return err
}

func EffectiveGroupRatio(c *gin.Context, userGroup, group string) types.GroupRatioInfo {
	result := types.GroupRatioInfo{GroupRatio: ratio_setting.GetGroupRatio(group), GroupSpecialRatio: -1}
	if ratio, ok := ratio_setting.GetGroupGroupRatio(userGroup, group); ok {
		result.GroupRatio, result.GroupSpecialRatio, result.HasSpecialRatio = ratio, ratio, true
	}
	if value, ok := c.Get(personalRatioContextKey); ok {
		if ratios, ok := value.(map[string]float64); ok {
			if ratio, ok := ratios[group]; ok {
				result.GroupRatio, result.GroupSpecialRatio, result.HasSpecialRatio = ratio, ratio, true
			}
		}
	}
	return result
}

func HasPersonalGroupRatio(c *gin.Context, group string) bool {
	value, _ := c.Get(personalRatioContextKey)
	ratios, _ := value.(map[string]float64)
	_, ok := ratios[group]
	return ok
}