fanqie-go/internal/fqapi/device_key_pool.go

package fqapi

import (
	"fanqie/internal/common"
	"fanqie/internal/config"
	"fmt"
	"sync"
	"time"
)

var (
	poolLogger = common.NewLogger("device_key_pool", "INFO", "")
	globalPool *DeviceKeyPool
	poolOnce   sync.Once
)

// DeviceKeyItem 设备密钥项
type DeviceKeyItem struct {
	ID           int       `json:"id"`
	DeviceHeader string    `json:"device_header"`
	Key          string    `json:"key"`
	KeyVer       int       `json:"key_ver"`
	KeyTs        int64     `json:"key_ts"` // 密钥注册时间戳（key_register_ts）
	IsValid      bool      `json:"is_valid"`
	CreatedAt    time.Time `json:"created_at"`
	LastUsedAt   time.Time `json:"last_used_at"`

	// 仅内存字段（不落库）
	KeyFetchedAt time.Time `json:"-"` // 密钥获取时间，超 24h 异步刷新
	KeyRefreshAt time.Time `json:"-"` // 上次刷新尝试时间，防刷新风暴
}

// DeviceKeyPool 设备密钥池
type DeviceKeyPool struct {
	mu           sync.RWMutex
	currentIndex int
	items        []*DeviceKeyItem
	enabled      bool
}

// GetDeviceKeyPool 获取全局设备密钥池实例
func GetDeviceKeyPool() *DeviceKeyPool {
	poolOnce.Do(func() {
		globalPool = &DeviceKeyPool{
			items:   make([]*DeviceKeyItem, 0),
			enabled: false,
		}
	})
	return globalPool
}

// Initialize 初始化设备密钥池
func (p *DeviceKeyPool) Initialize() error {
	cfg := config.GetConfig()
	if cfg == nil || !cfg.DeviceKeyPool.Enabled {
		poolLogger.Info("设备密钥池未启用")
		p.enabled = false
		return nil
	}

	p.enabled = true
	poolLogger.Info("开始初始化设备密钥池...")

	// 从数据库加载已有的设备
	if err := p.loadFromDB(); err != nil {
		poolLogger.Warning("从数据库加载设备失败: " + err.Error())
	}

	// 检查并补充设备到初始化数量
	validCount := p.getValidCount()
	initialSize := cfg.DeviceKeyPool.InitialSize

	if validCount < initialSize {
		needCount := initialSize - validCount
		poolLogger.Info(fmt.Sprintf("当前有效设备数: %d, 需要补充: %d", validCount, needCount))

		for i := 0; i < needCount; i++ {
			if err := p.addNewDevice(); err != nil {
				poolLogger.Error(fmt.Sprintf("添加设备失败: %v", err))
			}
		}
	}

	poolLogger.Info(fmt.Sprintf("设备密钥池初始化完成，当前有效设备数: %d", p.getValidCount()))

	// last_used_at 批量落库（GetNext 不再同步写库）
	p.StartDevicePoolFlusher()

	return nil
}

// IsEnabled 检查设备密钥池是否启用
func (p *DeviceKeyPool) IsEnabled() bool {
	return p.enabled
}

// GetNext 获取下一个可用的设备密钥（轮询方式）
func (p *DeviceKeyPool) GetNext() (*DeviceKeyItem, error) {
	if !p.enabled {
		return nil, fmt.Errorf("设备密钥池未启用")
	}

	p.mu.Lock()
	defer p.mu.Unlock()

	// 过滤有效设备
	validItems := make([]*DeviceKeyItem, 0)
	for _, item := range p.items {
		if item.IsValid {
			validItems = append(validItems, item)
		}
	}

	if len(validItems) == 0 {
		return nil, fmt.Errorf("没有可用的设备")
	}

	// 轮询获取下一个设备
	item := validItems[p.currentIndex%len(validItems)]
	p.currentIndex++

	// 只更新内存中的最后使用时间，由后台 flusher 定期批量落库
	// （此前每次取设备同步 UPDATE，SQLite 单写锁下是吞吐瓶颈）
	item.LastUsedAt = time.Now()

	// 密钥超 24h 未刷新：异步换新（本次请求仍用旧密钥，下次生效）
	if time.Since(item.KeyFetchedAt) > 20*time.Hour && time.Since(item.KeyRefreshAt) > 10*time.Minute {
		item.KeyRefreshAt = time.Now()
		header := item.DeviceHeader
		go func() {
			if ki, err := GetKey(header); err == nil {
				UpdateKeyInPool(header, ki)
			} else {
				poolLogger.Warning(fmt.Sprintf("刷新设备密钥失败: %v", err))
			}
		}()
	}

	// 异步检查是否需要补充
	go p.checkAndRefill()

	return item, nil
}

// MarkInvalid 标记设备为无效
func (p *DeviceKeyPool) MarkInvalid(deviceHeader string) error {
	if !p.enabled {
		return nil
	}

	p.mu.Lock()
	defer p.mu.Unlock()

	for _, item := range p.items {
		if item.DeviceHeader == deviceHeader {
			item.IsValid = false
			// poolLogger.Warning(fmt.Sprintf("设备 ID=%d 已标记为无效", item.ID))

			// 更新数据库
			if err := p.updateDeviceInDB(item); err != nil {
				return err
			}

			// 异步补充新设备
			go p.checkAndRefill()

			return nil
		}
	}

	return fmt.Errorf("未找到设备")
}

// RemoveInvalid 移除所有无效设备并自动补充
func (p *DeviceKeyPool) RemoveInvalid() error {
	if !p.enabled {
		return nil
	}

	p.mu.Lock()

	// 从数据库中删除无效设备
	if err := p.removeInvalidFromDB(); err != nil {
		p.mu.Unlock()
		return err
	}

	// 从内存中移除无效设备
	validItems := make([]*DeviceKeyItem, 0)
	for _, item := range p.items {
		if item.IsValid {
			validItems = append(validItems, item)
		}
	}

	removedCount := len(p.items) - len(validItems)
	p.items = validItems

	if removedCount > 0 {
		poolLogger.Info(fmt.Sprintf("已移除 %d 个无效设备", removedCount))

		// 立即检查并补充设备（使用已持有锁的版本）
		p.checkAndRefillLocked()
	}

	p.mu.Unlock()
	return nil
}

// GetStats 获取池统计信息
func (p *DeviceKeyPool) GetStats() map[string]interface{} {
	p.mu.RLock()
	defer p.mu.RUnlock()

	validCount := 0
	invalidCount := 0

	for _, item := range p.items {
		if item.IsValid {
			validCount++
		} else {
			invalidCount++
		}
	}

	return map[string]interface{}{
		"enabled":       p.enabled,
		"total":         len(p.items),
		"valid":         validCount,
		"invalid":       invalidCount,
		"current_index": p.currentIndex,
	}
}

// checkAndRefill 检查并补充设备（私有方法，会自动加锁）
func (p *DeviceKeyPool) checkAndRefill() {
	if !p.enabled {
		return
	}

	cfg := config.GetConfig()
	if cfg == nil {
		return
	}

	p.mu.Lock()
	validCount := p.getValidCount()
	p.mu.Unlock()

	minSize := cfg.DeviceKeyPool.MinSize
	maxSize := cfg.DeviceKeyPool.MaxSize

	if validCount < minSize {
		needCount := minSize - validCount
		// 限制不超过最大值
		if validCount+needCount > maxSize {
			needCount = maxSize - validCount
		}

		if needCount > 0 {
			poolLogger.Info(fmt.Sprintf("有效设备数 %d 低于最小值 %d，开始补充 %d 个设备", validCount, minSize, needCount))

			for i := 0; i < needCount; i++ {
				if err := p.addNewDevice(); err != nil {
					poolLogger.Error(fmt.Sprintf("补充设备失败: %v", err))
				}
			}
		}
	}
}

// checkAndRefillLocked 检查并补充设备（已持有锁的内部版本）
func (p *DeviceKeyPool) checkAndRefillLocked() {
	if !p.enabled {
		return
	}

	cfg := config.GetConfig()
	if cfg == nil {
		return
	}

	validCount := p.getValidCount()
	minSize := cfg.DeviceKeyPool.MinSize
	maxSize := cfg.DeviceKeyPool.MaxSize

	if validCount < minSize {
		needCount := minSize - validCount
		// 限制不超过最大值
		if validCount+needCount > maxSize {
			needCount = maxSize - validCount
		}

		if needCount > 0 {
			poolLogger.Info(fmt.Sprintf("有效设备数 %d 低于最小值 %d，开始补充 %d 个设备", validCount, minSize, needCount))

			// 临时释放锁，因为addNewDevice需要加锁
			p.mu.Unlock()
			for i := 0; i < needCount; i++ {
				if err := p.addNewDevice(); err != nil {
					poolLogger.Error(fmt.Sprintf("补充设备失败: %v", err))
				}
			}
			p.mu.Lock()
		}
	}
}

// getValidCount 获取有效设备数量（调用前需已加锁或在RLock内）
func (p *DeviceKeyPool) getValidCount() int {
	count := 0
	for _, item := range p.items {
		if item.IsValid {
			count++
		}
	}
	return count
}

// addNewDevice 添加新设备（会自动加锁）
func (p *DeviceKeyPool) addNewDevice() error {
	// 注册新设备
	deviceHeader, err := Register()
	if err != nil {
		return fmt.Errorf("注册设备失败: %w", err)
	}

	// 获取密钥
	keyInfo, err := GetKey(deviceHeader)
	if err != nil {
		return fmt.Errorf("获取密钥失败: %w", err)
	}

	// 创建设备项
	item := &DeviceKeyItem{
		DeviceHeader: deviceHeader,
		Key:          keyInfo.Key,
		KeyVer:       keyInfo.Ver,
		KeyTs:        keyInfo.Ts,
		IsValid:      true,
		CreatedAt:    time.Now(),
		LastUsedAt:   time.Now(),
		KeyFetchedAt: time.Now(),
	}

	// 保存到数据库
	if err := p.saveDeviceToDB(item); err != nil {
		return fmt.Errorf("保存设备到数据库失败: %w", err)
	}

	// 添加到内存
	p.mu.Lock()
	p.items = append(p.items, item)
	p.mu.Unlock()

	// poolLogger.Info(fmt.Sprintf("已添加新设备 ID=%d", item.ID))
	return nil
}

// ===== 数据库操作方法 =====

// loadFromDB 从数据库加载设备
func (p *DeviceKeyPool) loadFromDB() error {
	if cacheDB == nil {
		return fmt.Errorf("数据库未初始化")
	}

	rows, err := cacheDB.Query(`
		SELECT id, device_header, key, key_ver, COALESCE(key_ts, 0), is_valid, created_at, last_used_at
		FROM device_pool
		ORDER BY id
	`)
	if err != nil {
		return err
	}
	defer rows.Close()

	items := make([]*DeviceKeyItem, 0)
	for rows.Next() {
		var item DeviceKeyItem
		var createdAt, lastUsedAt int64

		err := rows.Scan(
			&item.ID,
			&item.DeviceHeader,
			&item.Key,
			&item.KeyVer,
			&item.KeyTs,
			&item.IsValid,
			&createdAt,
			&lastUsedAt,
		)
		if err != nil {
			poolLogger.Error("扫描行失败: " + err.Error())
			continue
		}

		item.CreatedAt = time.Unix(createdAt, 0)
		item.LastUsedAt = time.Unix(lastUsedAt, 0)
		item.KeyFetchedAt = time.Now() // 旧设备的密钥年龄未知，从加载起算
		items = append(items, &item)
	}

	p.mu.Lock()
	p.items = items
	p.mu.Unlock()

	poolLogger.Info(fmt.Sprintf("从数据库加载了 %d 个设备", len(items)))
	return nil
}

// boolToInt 布尔转整数存储（is_valid 列为 INTEGER，PostgreSQL 严格要求类型匹配）
func boolToInt(b bool) int {
	if b {
		return 1
	}
	return 0
}

// saveDeviceToDB 保存设备到数据库
func (p *DeviceKeyPool) saveDeviceToDB(item *DeviceKeyItem) error {
	if cacheDB == nil {
		return fmt.Errorf("数据库未初始化")
	}

	// RETURNING id 同时兼容 SQLite(3.35+) 与 PostgreSQL，替代 LastInsertId()
	err := cacheDB.QueryRow(q(`
		INSERT INTO device_pool (device_header, key, key_ver, key_ts, is_valid, created_at, last_used_at, updated_at)
		VALUES (?, ?, ?, ?, ?, ?, ?, ?)
		RETURNING id
	`), item.DeviceHeader, item.Key, item.KeyVer, item.KeyTs, boolToInt(item.IsValid),
		item.CreatedAt.Unix(), item.LastUsedAt.Unix(), item.CreatedAt.Unix()).Scan(&item.ID)

	return err
}

// updateDeviceInDB 更新设备信息
func (p *DeviceKeyPool) updateDeviceInDB(item *DeviceKeyItem) error {
	if cacheDB == nil {
		return fmt.Errorf("数据库未初始化")
	}

	_, err := cacheDB.Exec(q(`
		UPDATE device_pool
		SET device_header = ?, key = ?, key_ver = ?, key_ts = ?, is_valid = ?, last_used_at = ?, updated_at = ?
		WHERE id = ?
	`), item.DeviceHeader, item.Key, item.KeyVer, item.KeyTs, boolToInt(item.IsValid),
		item.LastUsedAt.Unix(), item.LastUsedAt.Unix(), item.ID)

	return err
}

// StartDevicePoolFlusher 启动 last_used_at 批量落库任务（GetNext 已改为只写内存）
func (p *DeviceKeyPool) StartDevicePoolFlusher() {
	go func() {
		for range time.Tick(time.Minute) {
			FlushDevicePool()
		}
	}()
	poolLogger.Info("设备池 last_used_at 批量落库已启动（每分钟）")
}

// FlushDevicePool 将内存中的设备使用时间批量刷入数据库
// ponytail: 进程退出时最多丢 1 分钟使用记录（仅统计用途），无需精确停表
func FlushDevicePool() {
	p := GetDeviceKeyPool()
	if !p.enabled || cacheDB == nil {
		return
	}
	p.mu.RLock()
	items := append([]*DeviceKeyItem(nil), p.items...)
	p.mu.RUnlock()
	for _, it := range items {
		p.updateDeviceInDB(it)
	}
}

// removeInvalidFromDB 从数据库移除无效设备
func (p *DeviceKeyPool) removeInvalidFromDB() error {
	if cacheDB == nil {
		return fmt.Errorf("数据库未初始化")
	}

	result, err := cacheDB.Exec(`DELETE FROM device_pool WHERE is_valid = 0`)
	if err != nil {
		return err
	}

	affected, _ := result.RowsAffected()
	if affected > 0 {
		poolLogger.Info(fmt.Sprintf("从数据库删除了 %d 个无效设备", affected))
	}

	return nil
}

// GetDeviceHeaderFromPool 从池中获取设备header（兼容旧接口）
func GetDeviceHeaderFromPool() string {
	pool := GetDeviceKeyPool()
	if !pool.IsEnabled() {
		return ""
	}

	item, err := pool.GetNext()
	if err != nil {
		poolLogger.Error("从池中获取设备失败: " + err.Error())
		return ""
	}

	return item.DeviceHeader
}

// GetKeyInfoFromPool 从池中获取密钥信息（兼容旧接口）
func GetKeyInfoFromPool(deviceHeader string) *KeyInfo {
	pool := GetDeviceKeyPool()
	if !pool.IsEnabled() {
		return nil
	}

	pool.mu.RLock()
	defer pool.mu.RUnlock()

	for _, item := range pool.items {
		if item.DeviceHeader == deviceHeader && item.IsValid {
			return &KeyInfo{
				Key: item.Key,
				Ver: item.KeyVer,
			}
		}
	}

	return nil
}

// UpdateKeyInPool 更新池中设备的密钥（含注册时间戳）
func UpdateKeyInPool(deviceHeader string, ki *KeyInfo) error {
	pool := GetDeviceKeyPool()
	if !pool.IsEnabled() || ki == nil {
		return nil
	}

	pool.mu.Lock()
	defer pool.mu.Unlock()

	for _, item := range pool.items {
		if item.DeviceHeader == deviceHeader {
			item.Key = ki.Key
			item.KeyVer = ki.Ver
			item.KeyTs = ki.Ts
			item.KeyFetchedAt = time.Now()
			return pool.updateDeviceInDB(item)
		}
	}

	return fmt.Errorf("未找到设备")
}

// GetAnyValidDeviceHeader 获取任意一个有效设备的header（用于辅助功能）
func GetAnyValidDeviceHeader() string {
	pool := GetDeviceKeyPool()
	if !pool.IsEnabled() {
		return ""
	}

	pool.mu.RLock()
	defer pool.mu.RUnlock()

	for _, item := range pool.items {
		if item.IsValid {
			return item.DeviceHeader
		}
	}

	return ""
}

// GetValidDevice 获取一个有效的设备
func GetValidDevice() (*DeviceKeyItem, error) {
	pool := GetDeviceKeyPool()
	if !pool.IsEnabled() {
		return nil, fmt.Errorf("设备池未启用")
	}
	return pool.GetNext()
}

// GetCurrentDevice 获取当前使用的设备（和GetValidDevice一样，返回下一个可用设备）
func GetCurrentDevice() *DeviceKeyItem {
	pool := GetDeviceKeyPool()
	if !pool.IsEnabled() {
		return nil
	}
	item, _ := pool.GetNext()
	return item
}