package service

import (
	"slices"

	"github.com/QuantumNous/new-api/model"
	relaycommon "github.com/QuantumNous/new-api/relay/common"
	"github.com/QuantumNous/new-api/relaykit/dto"
)

const (
	usageBillingPathLocal              = "local"
	usageBillingPathUpstream           = "upstream"
	usageBillingPathOpenAI             = "billing-usage-openai"
	usageBillingPathOpenAIEstimated    = "billing-usage-openai-estimated"
	usageBillingPathAnthropic          = "billing-usage-anthropic"
	usageBillingPathAnthropicEstimated = "billing-usage-anthropic-estimated"
	usageBillingPathGemini             = "billing-usage-gemini"
	usageBillingPathGeminiEstimated    = "billing-usage-gemini-estimated"
)

func effectiveBillingUsage(usage *dto.Usage) *dto.Usage {
	if billingUsage, ok := usageFromBillingUsage(usage); ok {
		return billingUsage
	}
	return usage
}

type claudeCacheReadAdjustment struct {
	Rule                string `json:"rule"`
	Model               string `json:"model"`
	OriginalInputTokens int    `json:"original_input_tokens"`
	InputTokens         int    `json:"input_tokens"`
	CacheReadTokens     int    `json:"cache_read_tokens"`
	Status              string `json:"status"`
}

// normalizeClaudeCacheReadUsage runs once after canonical upstream usage is
// selected. Only the billing copy changes; downstream responses and the
// original upstream snapshot retain their protocol values.
func normalizeClaudeCacheReadUsage(info *relaycommon.RelayInfo, usage *dto.Usage) (*dto.Usage, *claudeCacheReadAdjustment) {
	if info == nil || info.ChannelMeta == nil || usage == nil ||
		usage.UsageSemantic != dto.BillingUsageSemanticAnthropic ||
		usage.BillingUsage == nil || usage.BillingUsage.Estimated || usage.BillingUsage.ClaudeUsage == nil {
		return usage, nil
	}
	upstreamModel := info.GetUpstreamModelName()
	if info.ResponseModel != nil && info.ResponseModel.UpstreamModel != "" {
		// Some handlers replace UpstreamModelName with the response model.
		// Match the name sent upstream, never the provider's response declaration.
		upstreamModel = info.ResponseModel.UpstreamModel
	}
	if !slices.Contains(info.ChannelOtherSettings.ClaudeInputIncludesCacheReadModels, upstreamModel) {
		return usage, nil
	}
	cached := usage.PromptTokensDetails.CachedTokens
	adjustment := &claudeCacheReadAdjustment{
		Rule: "claude_input_includes_cache_read_models", Model: upstreamModel,
		OriginalInputTokens: usage.PromptTokens, InputTokens: usage.PromptTokens,
		CacheReadTokens: cached, Status: "invalid_usage",
	}
	if usage.PromptTokens < 0 || cached < 0 || cached > usage.PromptTokens ||
		usage.CompletionTokens < 0 || usage.InputTokens < cached || usage.TotalTokens < cached {
		return usage, adjustment
	}
	if cached == 0 {
		return usage, nil
	}
	adjusted := *usage
	adjusted.PromptTokens -= cached
	adjusted.InputTokens -= cached
	adjusted.TotalTokens -= cached
	adjusted.BillingUsage = dto.CloneBillingUsage(usage.BillingUsage)
	adjusted.BillingUsage.ClaudeUsage.InputTokens = adjusted.PromptTokens
	adjustment.InputTokens = adjusted.PromptTokens
	adjustment.Status = "applied"
	return &adjusted, adjustment
}

func usageBillingPathForLog(isLocalCountTokens bool, usage *dto.Usage) string {
	effectiveUsage, ok := usageFromBillingUsage(usage)
	if !ok {
		if isLocalCountTokens {
			return usageBillingPathLocal
		}
		return usageBillingPathUpstream
	}

	switch effectiveUsage.UsageSemantic {
	case dto.BillingUsageSemanticOpenAI:
		if usage.BillingUsage.Estimated {
			return usageBillingPathOpenAIEstimated
		}
		return usageBillingPathOpenAI
	case dto.BillingUsageSemanticAnthropic:
		if usage.BillingUsage.Estimated {
			return usageBillingPathAnthropicEstimated
		}
		return usageBillingPathAnthropic
	case dto.BillingUsageSemanticGemini:
		if usage.BillingUsage.Estimated {
			return usageBillingPathGeminiEstimated
		}
		return usageBillingPathGemini
	}

	return usageBillingPathUpstream
}

func appendUsageBillingPathForLog(other *model.LogOther, isLocalCountTokens bool, usage *dto.Usage) {
	if other == nil {
		return
	}
	other.SetAdmin("usage_billing_path", usageBillingPathForLog(isLocalCountTokens, usage))
}

func usageFromBillingUsage(usage *dto.Usage) (*dto.Usage, bool) {
	if usage == nil || usage.BillingUsage == nil {
		return nil, false
	}
	return usage.BillingUsage.CanonicalUsage()
}