package model

import (
	"context"
	"strings"
	"time"

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

type StatusChannel struct {
	Id     int
	Group  string
	Status int
}

type GroupCacheStats struct {
	Requests       int64 `json:"requests"`
	CachedRequests int64 `json:"cached_requests"`
	InputTokens    int64 `json:"input_tokens"`
	CachedTokens   int64 `json:"cached_tokens"`
	FirstResponseMs float64 `json:"first_response_ms,omitempty"`
}

func GetStatusChannels(ctx context.Context) ([]StatusChannel, error) {
	var channels []StatusChannel
	err := DB.WithContext(ctx).Model(&Channel{}).Select("id", "group", "status").
		Where("status IN ?", []int{common.ChannelStatusEnabled, common.ChannelStatusAutoDisabled}).
		Order("id").Find(&channels).Error
	return channels, err
}

// Channels retain their identity when a public group is renamed.
func GetGroupCacheStats(ctx context.Context, channels []StatusChannel, start, end time.Time) (map[string]GroupCacheStats, error) {
	stats := make(map[string]GroupCacheStats)
	firstResponseTotals := make(map[string]struct {
		total float64
		count int64
	})
	channelGroups := make(map[int][]string)
	ids := make([]int, 0, len(channels))
	for _, channel := range channels {
		ids = append(ids, channel.Id)
		for _, group := range strings.Split(channel.Group, ",") {
			group = strings.TrimSpace(group)
			if group != "" && !common.StringsContains(channelGroups[channel.Id], group) {
				channelGroups[channel.Id] = append(channelGroups[channel.Id], group)
				stats[group] = GroupCacheStats{}
			}
		}
	}
	if len(ids) == 0 {
		return stats, nil
	}
	rows, err := LOG_DB.WithContext(ctx).Model(&Log{}).Select("channel_id", "prompt_tokens", "other").
		Where("type = ? AND token_id > 0 AND created_at >= ? AND created_at < ?", LogTypeConsume, start.Unix(), end.Unix()).
		Where("channel_id IN ?", ids).Rows()
	if err != nil {
		return nil, err
	}
	defer rows.Close()
	for rows.Next() {
		var channelID int
		var prompt int64
		var raw string
		if err := rows.Scan(&channelID, &prompt, &raw); err != nil {
			return nil, err
		}
		var usage struct {
			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"`
			InputTotal int64  `json:"input_tokens_total"`
			Semantic   string `json:"usage_semantic"`
			Claude     bool   `json:"claude"`
			FirstResponseMs float64 `json:"frt"`
		}
		if strings.TrimSpace(raw) != "" && common.UnmarshalJsonStr(raw, &usage) != nil {
			continue
		}
		input := max(prompt, 0)
		cached := max(usage.Cached, 0)
		if usage.InputTotal > 0 {
			input = usage.InputTotal
		} else if usage.Semantic == "anthropic" || (usage.Semantic == "" && usage.Claude) {
			input += cached + max(usage.Written, usage.Created, max(usage.Created5m, 0)+max(usage.Created1h, 0), 0)
		}
		if input == 0 {
			continue
		}
		cached = min(cached, input)
		for _, group := range channelGroups[channelID] {
			total := stats[group]
			total.Requests++
			if cached > 0 {
				total.CachedRequests++
			}
			total.InputTokens += input
			total.CachedTokens += cached
			stats[group] = total
			if usage.FirstResponseMs > 0 {
				frt := firstResponseTotals[group]
				frt.total += usage.FirstResponseMs
				frt.count++
				firstResponseTotals[group] = frt
			}
		}
	}
	for group, total := range firstResponseTotals {
		if total.count > 0 {
			groupStats := stats[group]
			groupStats.FirstResponseMs = total.total / float64(total.count)
			stats[group] = groupStats
		}
	}
	return stats, rows.Err()
}