package model

import (
	"context"
	"time"

	"github.com/QuantumNous/new-api/common"
)

type GroupModelCost struct {
	USDPerMillion *float64 `json:"usd_per_million"`
	Requests      int64    `json:"requests"`
	Tokens        float64  `json:"tokens"`
	quota         float64
}

// Consumption logs contain the settled charge, not the pre-consumed quota.
// Attribute by the group actually billed, never by today's channel membership.
func GetGroupModelCosts(ctx context.Context, start, end time.Time) (map[string]map[string]GroupModelCost, error) {
	result := map[string]map[string]GroupModelCost{}
	rows, err := LOG_DB.WithContext(ctx).Model(&Log{}).Select("group", "model_name", "quota", "prompt_tokens", "completion_tokens", "other").
		Where("type = ? AND token_id > 0 AND created_at >= ? AND created_at < ?", LogTypeConsume, start.Unix(), end.Unix()).Rows()
	if err != nil {
		return nil, err
	}
	defer rows.Close()
	for rows.Next() {
		var group, name, raw string
		var quota, prompt, completion int64
		if err := rows.Scan(&group, &name, &quota, &prompt, &completion, &raw); err != nil {
			return nil, err
		}
		if group == "" || name == "" || quota < 0 {
			continue
		}
		tokens, valid := consumptionTokenTotal(prompt, completion, raw)
		if !valid || tokens <= 0 {
			continue
		}
		if result[group] == nil {
			result[group] = map[string]GroupModelCost{}
		}
		value := result[group][name]
		value.Requests++
		value.Tokens += tokens
		value.quota += float64(quota)
		result[group][name] = value
	}
	if err := rows.Err(); err != nil {
		return nil, err
	}
	for group, models := range result {
		for name, value := range models {
			if common.QuotaPerUnit > 0 && value.Tokens > 0 {
				rate := value.quota / common.QuotaPerUnit / value.Tokens * 1_000_000
				value.USDPerMillion = &rate
			}
			result[group][name] = value
		}
	}
	return result, nil
}

func consumptionTokenTotal(prompt, completion int64, raw string) (float64, bool) {
	var usage struct {
		Total     int64  `json:"input_tokens_total"`
		Cached    int64  `json:"cache_tokens"`
		Created   int64  `json:"cache_creation_tokens"`
		Created5m int64  `json:"cache_creation_tokens_5m"`
		Created1h int64  `json:"cache_creation_tokens_1h"`
		Written   int64  `json:"cache_write_tokens"`
		Semantic  string `json:"usage_semantic"`
		Claude    bool   `json:"claude"`
	}
	if raw != "" && common.UnmarshalJsonStr(raw, &usage) != nil {
		return 0, false
	}
	input := float64(max(prompt, 0))
	if usage.Total > 0 {
		input = float64(usage.Total)
	} else if usage.Semantic == "anthropic" || (usage.Semantic == "" && usage.Claude) {
		created := max(float64(usage.Written), float64(usage.Created), float64(max(usage.Created5m, 0))+float64(max(usage.Created1h, 0)), 0)
		input += float64(max(usage.Cached, 0)) + created
	}
	return input + float64(max(completion, 0)), true
}