CyberStrikeAI/internal/database/model_token_usage.go

package database

import (
	"database/sql"
	"encoding/json"
	"fmt"
	"math"
	"strconv"
	"strings"
	"time"

	"github.com/google/uuid"
	"go.uber.org/zap"
)

const modelTokenUsageEventType = "eino_usage_summary"

// ModelTokenUsage records one model-usage summary emitted by an Agent run.
type ModelTokenUsage struct {
	ID               string    `json:"id"`
	ProcessDetailID  string    `json:"processDetailId"`
	MessageID        string    `json:"messageId"`
	ConversationID   string    `json:"conversationId"`
	ProjectID        string    `json:"projectId,omitempty"`
	Source           string    `json:"source"`
	Orchestration    string    `json:"orchestration"`
	Reason           string    `json:"reason"`
	Model            string    `json:"model,omitempty"`
	ModelCalls       int64     `json:"modelCalls"`
	PromptTokens     int64     `json:"promptTokens"`
	CompletionTokens int64     `json:"completionTokens"`
	TotalTokens      int64     `json:"totalTokens"`
	CachedTokens     int64     `json:"cachedTokens"`
	ReasoningTokens  int64     `json:"reasoningTokens"`
	CreatedAt        time.Time `json:"createdAt"`
	UpdatedAt        time.Time `json:"updatedAt"`
}

// ModelTokenUsageSummary is the aggregate shape used by dashboard and APIs.
type ModelTokenUsageSummary struct {
	Events           int64 `json:"events"`
	ModelCalls       int64 `json:"modelCalls"`
	PromptTokens     int64 `json:"promptTokens"`
	CompletionTokens int64 `json:"completionTokens"`
	TotalTokens      int64 `json:"totalTokens"`
	CachedTokens     int64 `json:"cachedTokens"`
	ReasoningTokens  int64 `json:"reasoningTokens"`
}

// ModelTokenUsageBreakdown is a grouped aggregate row.
type ModelTokenUsageBreakdown struct {
	Key              string `json:"key"`
	Label            string `json:"label,omitempty"`
	Events           int64  `json:"events"`
	ModelCalls       int64  `json:"modelCalls"`
	PromptTokens     int64  `json:"promptTokens"`
	CompletionTokens int64  `json:"completionTokens"`
	TotalTokens      int64  `json:"totalTokens"`
	CachedTokens     int64  `json:"cachedTokens"`
	ReasoningTokens  int64  `json:"reasoningTokens"`
}

// ModelTokenUsageStats is a compact API response for usage dashboards.
type ModelTokenUsageStats struct {
	Summary         ModelTokenUsageSummary     `json:"summary"`
	Today           ModelTokenUsageSummary     `json:"today"`
	ByDay           []ModelTokenUsageBreakdown `json:"byDay"`
	ByModel         []ModelTokenUsageBreakdown `json:"byModel"`
	ByOrchestration []ModelTokenUsageBreakdown `json:"byOrchestration"`
	Recent          []ModelTokenUsage          `json:"recent"`
}

// ModelTokenUsageFilter scopes usage queries.
type ModelTokenUsageFilter struct {
	ConversationID string
	ProjectID      string
	Since          time.Time
	Until          time.Time
	Days           int
	Access         RBACListAccess
	Limit          int
}

func modelTokenUsageFromProcessDetail(messageID, conversationID, processDetailID string, data interface{}) (ModelTokenUsage, bool) {
	m := mapFromUsageData(data)
	if len(m) == 0 {
		return ModelTokenUsage{}, false
	}
	usage := ModelTokenUsage{
		ID:               uuid.New().String(),
		ProcessDetailID:  strings.TrimSpace(processDetailID),
		MessageID:        strings.TrimSpace(messageID),
		ConversationID:   strings.TrimSpace(conversationID),
		Source:           strings.TrimSpace(fmt.Sprint(m["source"])),
		Orchestration:    strings.TrimSpace(fmt.Sprint(m["orchestration"])),
		Reason:           strings.TrimSpace(fmt.Sprint(m["reason"])),
		Model:            strings.TrimSpace(fmt.Sprint(m["model"])),
		ModelCalls:       usageInt64(m["modelCalls"]),
		PromptTokens:     usageInt64(m["promptTokens"]),
		CompletionTokens: usageInt64(m["completionTokens"]),
		TotalTokens:      usageInt64(m["totalTokens"]),
		CachedTokens:     usageInt64(m["cachedTokens"]),
		ReasoningTokens:  usageInt64(m["reasoningTokens"]),
	}
	if usage.TotalTokens == 0 && (usage.PromptTokens > 0 || usage.CompletionTokens > 0) {
		usage.TotalTokens = usage.PromptTokens + usage.CompletionTokens
	}
	if usage.ProcessDetailID == "" || usage.MessageID == "" || usage.ConversationID == "" {
		return ModelTokenUsage{}, false
	}
	if usage.ModelCalls == 0 && usage.TotalTokens == 0 && usage.PromptTokens == 0 && usage.CompletionTokens == 0 && usage.CachedTokens == 0 && usage.ReasoningTokens == 0 {
		return ModelTokenUsage{}, false
	}
	return usage, true
}

func mapFromUsageData(data interface{}) map[string]interface{} {
	switch v := data.(type) {
	case nil:
		return nil
	case map[string]interface{}:
		return v
	case string:
		var m map[string]interface{}
		if err := json.Unmarshal([]byte(v), &m); err == nil {
			return m
		}
	case []byte:
		var m map[string]interface{}
		if err := json.Unmarshal(v, &m); err == nil {
			return m
		}
	default:
		raw, err := json.Marshal(v)
		if err == nil {
			var m map[string]interface{}
			if err := json.Unmarshal(raw, &m); err == nil {
				return m
			}
		}
	}
	return nil
}

func usageInt64(v interface{}) int64 {
	switch n := v.(type) {
	case int:
		return int64(n)
	case int8:
		return int64(n)
	case int16:
		return int64(n)
	case int32:
		return int64(n)
	case int64:
		return n
	case uint:
		return int64(n)
	case uint8:
		return int64(n)
	case uint16:
		return int64(n)
	case uint32:
		return int64(n)
	case uint64:
		if n > math.MaxInt64 {
			return math.MaxInt64
		}
		return int64(n)
	case float32:
		return int64(n)
	case float64:
		return int64(n)
	case json.Number:
		i, _ := n.Int64()
		return i
	case string:
		i, _ := strconv.ParseInt(strings.TrimSpace(n), 10, 64)
		return i
	default:
		i, _ := strconv.ParseInt(strings.TrimSpace(fmt.Sprint(v)), 10, 64)
		return i
	}
}

func (db *DB) maybeRecordModelTokenUsage(messageID, conversationID, processDetailID, eventType string, data interface{}) {
	if db == nil || eventType != modelTokenUsageEventType {
		return
	}
	usage, ok := modelTokenUsageFromProcessDetail(messageID, conversationID, processDetailID, data)
	if !ok {
		return
	}
	if err := db.UpsertModelTokenUsage(usage); err != nil && db.logger != nil {
		db.logger.Warn("保存模型Token用量失败",
			zap.String("processDetailId", processDetailID),
			zap.String("conversationId", conversationID),
			zap.Error(err))
	}
}

// UpsertModelTokenUsage persists usage with process_detail_id idempotency.
func (db *DB) UpsertModelTokenUsage(usage ModelTokenUsage) error {
	if db == nil {
		return fmt.Errorf("database is nil")
	}
	now := time.Now()
	createdAt := usage.CreatedAt
	if createdAt.IsZero() {
		createdAt = now
	}
	if usage.ID == "" {
		usage.ID = uuid.New().String()
	}
	var projectID sql.NullString
	if err := db.QueryRow(`SELECT project_id FROM conversations WHERE id = ?`, usage.ConversationID).Scan(&projectID); err != nil && err != sql.ErrNoRows {
		return fmt.Errorf("查询对话项目失败: %w", err)
	}
	projectValue := interface{}(nil)
	if projectID.Valid && strings.TrimSpace(projectID.String) != "" {
		projectValue = strings.TrimSpace(projectID.String)
	}
	_, err := db.Exec(`
INSERT INTO model_token_usage (
	id, process_detail_id, message_id, conversation_id, project_id,
	source, orchestration, reason, model, model_calls,
	prompt_tokens, completion_tokens, total_tokens, cached_tokens, reasoning_tokens,
	created_at, updated_at
) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
ON CONFLICT(process_detail_id) DO UPDATE SET
	message_id = excluded.message_id,
	conversation_id = excluded.conversation_id,
	project_id = excluded.project_id,
	source = excluded.source,
	orchestration = excluded.orchestration,
	reason = excluded.reason,
	model = excluded.model,
	model_calls = excluded.model_calls,
	prompt_tokens = excluded.prompt_tokens,
	completion_tokens = excluded.completion_tokens,
	total_tokens = excluded.total_tokens,
	cached_tokens = excluded.cached_tokens,
	reasoning_tokens = excluded.reasoning_tokens,
	created_at = excluded.created_at,
	updated_at = excluded.updated_at`,
		usage.ID, usage.ProcessDetailID, usage.MessageID, usage.ConversationID, projectValue,
		usage.Source, usage.Orchestration, usage.Reason, usage.Model, usage.ModelCalls,
		usage.PromptTokens, usage.CompletionTokens, usage.TotalTokens, usage.CachedTokens, usage.ReasoningTokens,
		createdAt, now,
	)
	if err != nil {
		return fmt.Errorf("写入模型Token用量失败: %w", err)
	}
	return nil
}

// BackfillModelTokenUsageFromProcessDetails makes existing timeline usage events queryable.
func (db *DB) BackfillModelTokenUsageFromProcessDetails() error {
	if db == nil {
		return nil
	}
	rows, err := db.Query(`
SELECT pd.id, pd.message_id, pd.conversation_id, pd.data, pd.created_at
FROM process_details pd
LEFT JOIN model_token_usage mtu ON mtu.process_detail_id = pd.id
WHERE pd.event_type = ?
	AND (mtu.id IS NULL OR mtu.created_at != pd.created_at)`, modelTokenUsageEventType)
	if err != nil {
		return fmt.Errorf("查询历史模型Token用量失败: %w", err)
	}
	defer rows.Close()
	for rows.Next() {
		var processDetailID, messageID, conversationID string
		var data sql.NullString
		var createdAt string
		if err := rows.Scan(&processDetailID, &messageID, &conversationID, &data, &createdAt); err != nil {
			return fmt.Errorf("扫描历史模型Token用量失败: %w", err)
		}
		if !data.Valid {
			continue
		}
		usage, ok := modelTokenUsageFromProcessDetail(messageID, conversationID, processDetailID, data.String)
		if !ok {
			continue
		}
		usage.CreatedAt = parseModelTokenUsageTime(createdAt)
		if err := db.UpsertModelTokenUsage(usage); err != nil {
			return err
		}
	}
	if err := rows.Err(); err != nil {
		return fmt.Errorf("遍历历史模型Token用量失败: %w", err)
	}
	return nil
}

func (db *DB) GetModelTokenUsageStats(filter ModelTokenUsageFilter) (*ModelTokenUsageStats, error) {
	if db == nil {
		return nil, fmt.Errorf("database is nil")
	}
	if filter.Days <= 0 {
		filter.Days = 7
	}
	if filter.Limit <= 0 {
		filter.Limit = 10
	}
	where, args := buildModelTokenUsageWhere(filter, "mtu", "c")
	summary, err := db.queryModelTokenUsageSummary("SELECT "+modelTokenUsageSummarySelect("mtu")+" FROM model_token_usage mtu JOIN conversations c ON c.id = mtu.conversation_id"+where, args...)
	if err != nil {
		return nil, err
	}
	todayFilter := filter
	now := time.Now()
	todayFilter.Since = time.Date(now.Year(), now.Month(), now.Day(), 0, 0, 0, 0, now.Location())
	todayWhere, todayArgs := buildModelTokenUsageWhere(todayFilter, "mtu", "c")
	today, err := db.queryModelTokenUsageSummary("SELECT "+modelTokenUsageSummarySelect("mtu")+" FROM model_token_usage mtu JOIN conversations c ON c.id = mtu.conversation_id"+todayWhere, todayArgs...)
	if err != nil {
		return nil, err
	}
	byDay, err := db.queryModelTokenUsageBreakdown(
		"SELECT date(mtu.created_at) AS k, date(mtu.created_at) AS label, "+modelTokenUsageSummarySelect("mtu")+" FROM model_token_usage mtu JOIN conversations c ON c.id = mtu.conversation_id"+where+" GROUP BY date(mtu.created_at) ORDER BY k DESC LIMIT ?",
		append(args, filter.Days)...,
	)
	if err != nil {
		return nil, err
	}
	byModel, err := db.queryModelTokenUsageBreakdown(
		"SELECT COALESCE(NULLIF(TRIM(mtu.model), ''), 'unknown') AS k, COALESCE(NULLIF(TRIM(mtu.model), ''), 'Unknown') AS label, "+modelTokenUsageSummarySelect("mtu")+" FROM model_token_usage mtu JOIN conversations c ON c.id = mtu.conversation_id"+where+" GROUP BY k ORDER BY SUM(mtu.total_tokens) DESC LIMIT ?",
		append(args, filter.Limit)...,
	)
	if err != nil {
		return nil, err
	}
	byOrch, err := db.queryModelTokenUsageBreakdown(
		"SELECT COALESCE(NULLIF(TRIM(mtu.orchestration), ''), 'unknown') AS k, COALESCE(NULLIF(TRIM(mtu.orchestration), ''), 'Unknown') AS label, "+modelTokenUsageSummarySelect("mtu")+" FROM model_token_usage mtu JOIN conversations c ON c.id = mtu.conversation_id"+where+" GROUP BY k ORDER BY SUM(mtu.total_tokens) DESC LIMIT ?",
		append(args, filter.Limit)...,
	)
	if err != nil {
		return nil, err
	}
	recent, err := db.ListModelTokenUsage(filter)
	if err != nil {
		return nil, err
	}
	return &ModelTokenUsageStats{
		Summary:         summary,
		Today:           today,
		ByDay:           byDay,
		ByModel:         byModel,
		ByOrchestration: byOrch,
		Recent:          recent,
	}, nil
}

func modelTokenUsageSummarySelect(alias string) string {
	p := ""
	if alias != "" {
		p = alias + "."
	}
	return fmt.Sprintf(`COUNT(%sid),
COALESCE(SUM(%smodel_calls), 0),
COALESCE(SUM(%sprompt_tokens), 0),
COALESCE(SUM(%scompletion_tokens), 0),
COALESCE(SUM(%stotal_tokens), 0),
COALESCE(SUM(%scached_tokens), 0),
COALESCE(SUM(%sreasoning_tokens), 0)`, p, p, p, p, p, p, p)
}

func buildModelTokenUsageWhere(filter ModelTokenUsageFilter, usageAlias, convAlias string) (string, []interface{}) {
	where := " WHERE 1=1"
	args := []interface{}{}
	uPrefix := ""
	if usageAlias != "" {
		uPrefix = usageAlias + "."
	}
	if cid := strings.TrimSpace(filter.ConversationID); cid != "" {
		where += " AND " + uPrefix + "conversation_id = ?"
		args = append(args, cid)
	}
	where, args = appendConversationProjectFilter(where, args, filter.ProjectID, usageAlias)
	if !filter.Since.IsZero() {
		where += " AND " + uPrefix + "created_at >= ?"
		args = append(args, filter.Since)
	}
	if !filter.Until.IsZero() {
		where += " AND " + uPrefix + "created_at <= ?"
		args = append(args, filter.Until)
	}
	where, args = appendConversationAccessFilter(where, args, filter.Access.UserID, filter.Access.Scope, convAlias)
	return where, args
}

func (db *DB) queryModelTokenUsageSummary(query string, args ...interface{}) (ModelTokenUsageSummary, error) {
	var s ModelTokenUsageSummary
	err := db.QueryRow(query, args...).Scan(
		&s.Events, &s.ModelCalls, &s.PromptTokens, &s.CompletionTokens,
		&s.TotalTokens, &s.CachedTokens, &s.ReasoningTokens,
	)
	if err != nil {
		return s, fmt.Errorf("查询模型Token用量汇总失败: %w", err)
	}
	return s, nil
}

func (db *DB) queryModelTokenUsageBreakdown(query string, args ...interface{}) ([]ModelTokenUsageBreakdown, error) {
	rows, err := db.Query(query, args...)
	if err != nil {
		return nil, fmt.Errorf("查询模型Token用量分组失败: %w", err)
	}
	defer rows.Close()
	out := []ModelTokenUsageBreakdown{}
	for rows.Next() {
		var row ModelTokenUsageBreakdown
		if err := rows.Scan(
			&row.Key, &row.Label, &row.Events, &row.ModelCalls, &row.PromptTokens,
			&row.CompletionTokens, &row.TotalTokens, &row.CachedTokens, &row.ReasoningTokens,
		); err != nil {
			return nil, fmt.Errorf("扫描模型Token用量分组失败: %w", err)
		}
		out = append(out, row)
	}
	if err := rows.Err(); err != nil {
		return nil, fmt.Errorf("遍历模型Token用量分组失败: %w", err)
	}
	return out, nil
}

func (db *DB) ListModelTokenUsage(filter ModelTokenUsageFilter) ([]ModelTokenUsage, error) {
	if filter.Limit <= 0 {
		filter.Limit = 20
	}
	if filter.Limit > 500 {
		filter.Limit = 500
	}
	where, args := buildModelTokenUsageWhere(filter, "mtu", "c")
	args = append(args, filter.Limit)
	rows, err := db.Query(`
SELECT mtu.id, mtu.process_detail_id, mtu.message_id, mtu.conversation_id,
	COALESCE(mtu.project_id, ''), mtu.source, mtu.orchestration, mtu.reason, mtu.model,
	mtu.model_calls, mtu.prompt_tokens, mtu.completion_tokens, mtu.total_tokens,
	mtu.cached_tokens, mtu.reasoning_tokens, mtu.created_at, mtu.updated_at
FROM model_token_usage mtu
JOIN conversations c ON c.id = mtu.conversation_id`+where+`
ORDER BY mtu.created_at DESC, mtu.rowid DESC
LIMIT ?`, args...)
	if err != nil {
		return nil, fmt.Errorf("查询模型Token用量明细失败: %w", err)
	}
	defer rows.Close()
	out := []ModelTokenUsage{}
	for rows.Next() {
		var u ModelTokenUsage
		var createdAt, updatedAt string
		if err := rows.Scan(
			&u.ID, &u.ProcessDetailID, &u.MessageID, &u.ConversationID, &u.ProjectID,
			&u.Source, &u.Orchestration, &u.Reason, &u.Model, &u.ModelCalls,
			&u.PromptTokens, &u.CompletionTokens, &u.TotalTokens, &u.CachedTokens,
			&u.ReasoningTokens, &createdAt, &updatedAt,
		); err != nil {
			return nil, fmt.Errorf("扫描模型Token用量明细失败: %w", err)
		}
		u.CreatedAt = parseModelTokenUsageTime(createdAt)
		u.UpdatedAt = parseModelTokenUsageTime(updatedAt)
		out = append(out, u)
	}
	if err := rows.Err(); err != nil {
		return nil, fmt.Errorf("遍历模型Token用量明细失败: %w", err)
	}
	return out, nil
}

func parseModelTokenUsageTime(s string) time.Time {
	for _, layout := range []string{
		"2006-01-02 15:04:05.999999999-07:00",
		"2006-01-02 15:04:05.999999-07:00",
		"2006-01-02 15:04:05",
		time.RFC3339Nano,
		time.RFC3339,
	} {
		if t, err := time.Parse(layout, strings.TrimSpace(s)); err == nil {
			return t
		}
	}
	return time.Time{}
}