package service

import (
	"errors"
	"fmt"
	"math"

	"github.com/QuantumNous/new-api/common"
	"github.com/QuantumNous/new-api/model"
	"github.com/QuantumNous/new-api/pkg/billingexpr"
	benchmark "github.com/QuantumNous/new-api/pkg/model_benchmark"
	fingerprint "github.com/QuantumNous/new-api/pkg/model_fingerprint"
	"github.com/QuantumNous/new-api/relaykit/dto"
	"github.com/QuantumNous/new-api/setting/billing_setting"
)

type DetectionPrice struct {
	Snapshot billingexpr.BillingSnapshot `json:"snapshot"`
	Request  billingexpr.RequestInput    `json:"request"`
}

func DetectionPricing(userGroup, group, kind string) (DetectionPrice, int, error) {
	expression, ok := billing_setting.GetBillingExprForGroup(group, fingerprint.Model)
	if !ok {
		snapshot, err := model.GetModelPricingSnapshot([]string{fingerprint.Model})
		if err != nil || snapshot == nil || len(snapshot.Entries) != 1 {
			return DetectionPrice{}, 0, errors.New("pricing unavailable")
		}
		converted, err := model.PreviewModelPricingConversion(fingerprint.Model, snapshot.Entries[0].Configured)
		if err != nil || converted == nil || converted.Expression == "" {
			return DetectionPrice{}, 0, errors.New("unsupported detection pricing")
		}
		expression = converted.Expression
	}
	ratio := GetUserGroupRatio(userGroup, group)
	if ratio < 0 || math.IsNaN(ratio) || math.IsInf(ratio, 0) {
		return DetectionPrice{}, 0, errors.New("invalid group ratio")
	}
	price := DetectionPrice{Snapshot: billingexpr.BillingSnapshot{BillingMode: "tiered_expr", ModelName: fingerprint.Model, ExprString: expression, ExprHash: billingexpr.ExprHashString(expression), GroupRatio: ratio, QuotaPerUnit: common.QuotaPerUnit, ExprVersion: 1}}
	effort := "low"
	output, calls := fingerprint.OutputTokens, fingerprint.ProbeCount
	if kind == "benchmark" {
		output, calls = benchmark.OutputTokens, benchmark.ProbeCount
	} else if kind == "vision" {
		output, calls = VisionOutputTokens, 1
		effort = "medium"
	} else if kind != "fingerprint" {
		return price, 0, errors.New("invalid detection type")
	}
	price.Request.Body = []byte(fmt.Sprintf(`{"model":"gpt-6-astra","reasoning":{"effort":"%s"},"max_output_tokens":%d}`, effort, output))
	cost, err := billingexpr.DetectionCostBound(expression, 32768, float64(output))
	if err != nil {
		return price, 0, err
	}
	upper, upperClamp := common.QuotaRoundChecked(math.Ceil(cost / 1_000_000 * common.QuotaPerUnit * ratio))
	if upperClamp != nil {
		return price, 0, errors.New("reservation exceeds quota limit")
	}
	reserve, clamp := common.QuotaFromFloatChecked(float64(upper) * float64(calls))
	if clamp != nil {
		return price, 0, errors.New("reservation exceeds quota limit")
	}
	return price, reserve, nil
}

func DetectionUsageQuota(priceJSON string, input, output, cached int64, cacheWrites ...int64) (int, error) {
	written := int64(0)
	if len(cacheWrites) > 0 {
		written = cacheWrites[0]
	}
	if input < 0 || output < 0 || cached < 0 || written < 0 || cached > input || written > input-cached || input > 32768 || output > 16384 {
		return 0, errors.New("invalid upstream usage")
	}
	var price DetectionPrice
	if err := common.UnmarshalJsonStr(priceJSON, &price); err != nil {
		return 0, err
	}
	usage := &dto.Usage{PromptTokens: int(input), CompletionTokens: int(output)}
	usage.PromptTokensDetails.CachedTokens = int(cached)
	usage.PromptTokensDetails.CachedCreationTokens = int(written)
	params := BuildTieredTokenParams(usage, false, billingexpr.UsedVarsByHash(price.Snapshot.ExprString, price.Snapshot.ExprHash))
	result, err := billingexpr.ComputeTieredQuotaWithRequest(&price.Snapshot, params, price.Request)
	if err != nil || result.Clamp != nil || result.ActualQuotaAfterGroup < 0 {
		return 0, errors.New("usage settlement failed")
	}
	return result.ActualQuotaAfterGroup, nil
}