AIRouter/internal/usage/tracker.go

package usage

import (
	"database/sql"
	"fmt"
	"sync"
	"time"

	_ "github.com/mattn/go-sqlite3"
)

// AppUsage 记录单个应用的用量统计。
type AppUsage struct {
	AppID        string `json:"app_id"`
	RequestCount int64  `json:"request_count"`
	InputTokens  int64  `json:"input_tokens"`
	OutputTokens int64  `json:"output_tokens"`
	TotalTokens  int64  `json:"total_tokens"`
	ErrorCount   int64  `json:"error_count"`
	WindowStart  string `json:"window_start"`
	WindowEnd    string `json:"window_end"`
}

// Tracker 管理 per-app 用量统计（内存滑动窗口 + SQLite 持久化）。
type Tracker struct {
	mu            sync.RWMutex
	apps          map[string]*appCounters
	window        time.Duration
	db            *sql.DB
	logCh         chan usageLogEntry // 异步写入 channel
	stopCh        chan struct{}
	retentionDays int // 历史数据保留天数
}

type appCounters struct {
	requestCount int64
	inputTokens  int64
	outputTokens int64
	errorCount   int64
	windowStart  time.Time
}

type usageLogEntry struct {
	appID        string
	inputTokens  int
	outputTokens int
	isError      bool
	recordedAt   time.Time
}

// NewTracker 创建一个用量统计器。
// window 为统计窗口时长（如 1 小时），窗口到期后自动重置。
// dbPath 为 SQLite 持久化路径，空字符串则仅用内存。
// retentionDays 为历史数据保留天数，超过此天数的数据自动清理（0 表示不清理）。
func NewTracker(window time.Duration, dbPath string, retentionDays int) (*Tracker, error) {
	t := &Tracker{
		apps:          make(map[string]*appCounters),
		window:        window,
		logCh:         make(chan usageLogEntry, 1024),
		stopCh:        make(chan struct{}),
		retentionDays: retentionDays,
	}

	if dbPath != "" {
		db, err := sql.Open("sqlite3", dbPath)
		if err != nil {
			return nil, fmt.Errorf("open usage db: %w", err)
		}
		if err := initUsageDB(db); err != nil {
			return nil, fmt.Errorf("init usage db: %w", err)
		}
		t.db = db
		go t.logWriter() // 启动异步写入 goroutine
	}

	return t, nil
}

func initUsageDB(db *sql.DB) error {
	schema := `
	CREATE TABLE IF NOT EXISTS usage_logs (
		id INTEGER PRIMARY KEY AUTOINCREMENT,
		app_id TEXT NOT NULL,
		input_tokens INTEGER NOT NULL DEFAULT 0,
		output_tokens INTEGER NOT NULL DEFAULT 0,
		is_error INTEGER NOT NULL DEFAULT 0,
		recorded_at TEXT NOT NULL
	);
	CREATE INDEX IF NOT EXISTS idx_usage_app ON usage_logs(app_id);
	CREATE INDEX IF NOT EXISTS idx_usage_time ON usage_logs(recorded_at);`
	_, err := db.Exec(schema)
	return err
}

// Record 记录一次请求的用量。
func (t *Tracker) Record(appID string, inputTokens, outputTokens int, isError bool) {
	now := time.Now()

	// 更新内存计数器
	t.mu.Lock()
	c, ok := t.apps[appID]
	if !ok {
		c = &appCounters{windowStart: now}
		t.apps[appID] = c
	}

	// 窗口过期则重置
	if now.Sub(c.windowStart) > t.window {
		c.requestCount = 0
		c.inputTokens = 0
		c.outputTokens = 0
		c.errorCount = 0
		c.windowStart = now
	}

	c.requestCount++
	c.inputTokens += int64(inputTokens)
	c.outputTokens += int64(outputTokens)
	if isError {
		c.errorCount++
	}
	t.mu.Unlock()

	// 异步写入 SQLite（非阻塞，channel 满则丢弃）
	if t.db != nil {
		select {
		case t.logCh <- usageLogEntry{appID, inputTokens, outputTokens, isError, now}:
		default: // channel 满则丢弃，避免阻塞主流程
		}
	}
}

// logWriter 后台 goroutine，从 channel 读取用量日志并批量写入 SQLite。
func (t *Tracker) logWriter() {
	batch := make([]usageLogEntry, 0, 64)
	ticker := time.NewTicker(5 * time.Second)
	defer ticker.Stop()

	flush := func() {
		if len(batch) == 0 {
			return
		}
		tx, err := t.db.Begin()
		if err != nil {
			batch = batch[:0]
			return
		}
		stmt, _ := tx.Prepare(`INSERT INTO usage_logs (app_id, input_tokens, output_tokens, is_error, recorded_at) VALUES (?, ?, ?, ?, ?)`)
		if stmt != nil {
			for _, e := range batch {
				errInt := 0
				if e.isError {
					errInt = 1
				}
				stmt.Exec(e.appID, e.inputTokens, e.outputTokens, errInt, e.recordedAt.Format(time.RFC3339))
			}
			stmt.Close()
		}
		tx.Commit()
		batch = batch[:0]
	}

	for {
		select {
		case e := <-t.logCh:
			batch = append(batch, e)
			if len(batch) >= 64 {
				flush()
			}
		case <-ticker.C:
			flush()
			// 定期清理过期数据
			if t.retentionDays > 0 {
				t.cleanupOldLogs()
			}
		case <-t.stopCh:
			// 排空 channel 后 flush
			for len(t.logCh) > 0 {
				batch = append(batch, <-t.logCh)
			}
			flush()
			return
		}
	}
}

// Get 返回指定应用的当前用量快照。
func (t *Tracker) Get(appID string) AppUsage {
	t.mu.RLock()
	defer t.mu.RUnlock()

	c, ok := t.apps[appID]
	if !ok {
		return AppUsage{AppID: appID}
	}
	return AppUsage{
		AppID:        appID,
		RequestCount: c.requestCount,
		InputTokens:  c.inputTokens,
		OutputTokens: c.outputTokens,
		TotalTokens:  c.inputTokens + c.outputTokens,
		ErrorCount:   c.errorCount,
		WindowStart:  c.windowStart.Format(time.RFC3339),
		WindowEnd:    c.windowStart.Add(t.window).Format(time.RFC3339),
	}
}

// GetAll 返回所有应用的用量快照。
func (t *Tracker) GetAll() []AppUsage {
	t.mu.RLock()
	defer t.mu.RUnlock()

	result := make([]AppUsage, 0, len(t.apps))
	for appID, c := range t.apps {
		result = append(result, AppUsage{
			AppID:        appID,
			RequestCount: c.requestCount,
			InputTokens:  c.inputTokens,
			OutputTokens: c.outputTokens,
			TotalTokens:  c.inputTokens + c.outputTokens,
			ErrorCount:   c.errorCount,
			WindowStart:  c.windowStart.Format(time.RFC3339),
			WindowEnd:    c.windowStart.Add(t.window).Format(time.RFC3339),
		})
	}
	return result
}

// GetHistory 返回指定应用的历史用量（从 SQLite 查询，按小时聚合）。
// hours 为查询的时间范围（最近 N 小时）。
func (t *Tracker) GetHistory(appID string, hours int) ([]map[string]any, error) {
	if t.db == nil {
		return nil, fmt.Errorf("persistence not enabled")
	}
	if hours <= 0 {
		hours = 24
	}

	since := time.Now().Add(-time.Duration(hours) * time.Hour).Format(time.RFC3339)
	rows, err := t.db.Query(
		`SELECT 
			substr(recorded_at, 1, 13) as hour_bucket,
			COUNT(*) as request_count,
			SUM(input_tokens) as input_tokens,
			SUM(output_tokens) as output_tokens,
			SUM(is_error) as error_count
		FROM usage_logs 
		WHERE app_id = ? AND recorded_at >= ?
		GROUP BY hour_bucket
		ORDER BY hour_bucket DESC`,
		appID, since,
	)
	if err != nil {
		return nil, fmt.Errorf("query usage history: %w", err)
	}
	defer rows.Close()

	var result []map[string]any
	for rows.Next() {
		var hourBucket string
		var reqCount, inputTokens, outputTokens, errCount int64
		if err := rows.Scan(&hourBucket, &reqCount, &inputTokens, &outputTokens, &errCount); err != nil {
			return nil, err
		}
		result = append(result, map[string]any{
			"hour":          hourBucket,
			"request_count": reqCount,
			"input_tokens":  inputTokens,
			"output_tokens": outputTokens,
			"total_tokens":  inputTokens + outputTokens,
			"error_count":   errCount,
		})
	}
	return result, rows.Err()
}

// Close 关闭数据库连接。
func (t *Tracker) Close() error {
	close(t.stopCh)
	if t.db != nil {
		return t.db.Close()
	}
	return nil
}

// cleanupOldLogs 清理超过保留期限的用量日志。
func (t *Tracker) cleanupOldLogs() {
	cutoff := time.Now().Add(-time.Duration(t.retentionDays) * 24 * time.Hour).Format(time.RFC3339)
	if _, err := t.db.Exec(`DELETE FROM usage_logs WHERE recorded_at < ?`, cutoff); err != nil {
		// 清理失败不影响主流程，下次再试
		return
	}
}