package model

import (
	"context"
	"math"

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

type ModelPerformanceMetric struct {
	Value     float64 `json:"value"`
	Samples   int     `json:"samples"`
	Source    string  `json:"source"`
	UpdatedAt int64   `json:"updated_at"`
}

type GroupModelPerformance struct {
	Throughput *ModelPerformanceMetric `json:"throughput,omitempty"`
	FirstToken *ModelPerformanceMetric `json:"first_token,omitempty"`
}

type performanceAccumulator struct {
	total  float64
	count  int
	latest int64
}

func (a *performanceAccumulator) add(value *float64, at int64) {
	if value == nil {
		return
	}
	a.total += *value
	a.count++
	a.latest = max(a.latest, at)
}

func (a performanceAccumulator) metric(source string) *ModelPerformanceMetric {
	if a.count == 0 {
		return nil
	}
	return &ModelPerformanceMetric{Value: a.total / float64(a.count), Samples: a.count, Source: source, UpdatedAt: a.latest}
}

// The throughput definition matches usage logs: output tokens / total request
// seconds. New logs retain milliseconds; legacy logs use their stored seconds.
func modelLogPerformance(log Log) (throughput, firstToken *float64) {
	var other struct {
		FirstToken   *float64 `json:"frt"`
		LatencyMs    *float64 `json:"latency_ms"`
		StreamStatus struct {
			Status         string `json:"status"`
			EndReason      string `json:"end_reason"`
			ResponseStatus string `json:"response_status"`
		} `json:"stream_status"`
	}
	if common.UnmarshalJsonStr(log.Other, &other) != nil ||
		(other.StreamStatus.Status != "" && other.StreamStatus.Status != "ok") ||
		other.StreamStatus.EndReason == "client_gone" ||
		(other.StreamStatus.ResponseStatus != "" && other.StreamStatus.ResponseStatus != "completed") {
		return nil, nil
	}
	if log.IsStream && other.FirstToken != nil && *other.FirstToken >= 0 && !math.IsInf(*other.FirstToken, 0) && !math.IsNaN(*other.FirstToken) {
		firstToken = other.FirstToken
	}
	duration := float64(log.UseTime) * 1000
	if other.LatencyMs != nil {
		duration = *other.LatencyMs
	}
	if log.CompletionTokens > 0 && duration > 0 && !math.IsInf(duration, 0) && !math.IsNaN(duration) {
		value := float64(log.CompletionTokens) * 1000 / duration
		if value > 0 && !math.IsInf(value, 0) {
			throughput = &value
		}
	}
	return throughput, firstToken
}

// Attribute calls to their billed group, never to a channel's current groups.
// Each metric keeps its latest thirty usable calls independently. Failed and
// non-streaming calls cannot replace a known first-token measurement.
func GetGroupModelPerformance(ctx context.Context, group, modelName string) (GroupModelPerformance, error) {
	result := GroupModelPerformance{}
	if LOG_DB == nil || group == "" || modelName == "" {
		return result, nil
	}
	var throughput, firstToken performanceAccumulator
	lastID, lastAt := int64(0), int64(0)
	for {
		var logs []Log
		query := LOG_DB.WithContext(ctx).Model(&Log{}).
			Select("id", "created_at", "completion_tokens", "use_time", "is_stream", "other").
			Where(map[string]any{"type": LogTypeConsume, "group": group, "model_name": modelName}).
			Where("token_id > 0").Order("created_at DESC, id DESC").Limit(30)
		if throughput.count == 30 {
			query = query.Where("is_stream = ?", true)
		}
		if lastID > 0 {
			query = query.Where("created_at < ? OR (created_at = ? AND id < ?)", lastAt, lastAt, lastID)
		}
		if err := query.Find(&logs).Error; err != nil {
			return result, err
		}
		for _, log := range logs {
			tps, ttft := modelLogPerformance(log)
			if throughput.count < 30 {
				throughput.add(tps, log.CreatedAt)
			}
			if firstToken.count < 30 {
				firstToken.add(ttft, log.CreatedAt)
			}
		}
		if len(logs) < 30 || (throughput.count == 30 && firstToken.count == 30) {
			break
		}
		lastID = int64(logs[len(logs)-1].Id)
		lastAt = logs[len(logs)-1].CreatedAt
	}
	result.Throughput = throughput.metric("calls")
	result.FirstToken = firstToken.metric("calls")
	if result.Throughput != nil && result.FirstToken != nil {
		return result, nil
	}
	var channelIDs []int
	if err := DB.WithContext(ctx).Model(&Ability{}).
		Where(map[string]any{"group": group, "model": modelName, "enabled": true}).
		Pluck("channel_id", &channelIDs).Error; err != nil {
		return result, err
	}
	if len(channelIDs) == 0 {
		return result, nil
	}
	if err := DB.WithContext(ctx).Model(&Channel{}).
		Where("id IN ? AND status = ?", channelIDs, common.ChannelStatusEnabled).
		Pluck("id", &channelIDs).Error; err != nil {
		return result, err
	}
	if len(channelIDs) == 0 {
		return result, nil
	}
	rows, err := LOG_DB.WithContext(ctx).Model(&Log{}).
		Select("id", "channel_id", "created_at", "completion_tokens", "use_time", "is_stream", "other").
		Where("type = ? AND token_id = 0 AND model_name = ? AND is_stream = ? AND channel_id IN ?", LogTypeConsume, modelName, true, channelIDs).
		Where("content IN ?", []string{"模型测试", "Model test"}).Order("created_at DESC, id DESC").Rows()
	if err != nil {
		return result, err
	}
	defer rows.Close()
	var testThroughput, testFirstToken performanceAccumulator
	tpsChannels, ttftChannels := map[int]bool{}, map[int]bool{}
	for rows.Next() {
		var log Log
		if err := LOG_DB.ScanRows(rows, &log); err != nil {
			return result, err
		}
		tps, ttft := modelLogPerformance(log)
		if tps != nil && !tpsChannels[log.ChannelId] {
			testThroughput.add(tps, log.CreatedAt)
			tpsChannels[log.ChannelId] = true
		}
		if ttft != nil && !ttftChannels[log.ChannelId] {
			testFirstToken.add(ttft, log.CreatedAt)
			ttftChannels[log.ChannelId] = true
		}
		if (result.Throughput != nil || len(tpsChannels) == len(channelIDs)) &&
			(result.FirstToken != nil || len(ttftChannels) == len(channelIDs)) {
			break
		}
	}
	if err := rows.Err(); err != nil {
		return result, err
	}
	if result.Throughput == nil {
		result.Throughput = testThroughput.metric("test")
	}
	if result.FirstToken == nil {
		result.FirstToken = testFirstToken.metric("test")
	}
	return result, nil
}