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"`
	IsValid      bool      `json:"is_valid"`
	CreatedAt    time.Time `json:"created_at"`
	LastUsedAt   time.Time `json:"last_used_at"`
}

// 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()))
	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++
	
	// 更新最后使用时间
	item.LastUsedAt = time.Now()
	p.updateDeviceInDB(item)

	// poolLogger.Debug(fmt.Sprintf("使用设备 ID=%d", item.ID))
	
	// 异步检查是否需要补充
	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,
		IsValid:      true,
		CreatedAt:    time.Now(),
		LastUsedAt:   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, 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.IsValid,
			&createdAt,
			&lastUsedAt,
		)
		if err != nil {
			poolLogger.Error("扫描行失败: " + err.Error())
			continue
		}
		
		item.CreatedAt = time.Unix(createdAt, 0)
		item.LastUsedAt = time.Unix(lastUsedAt, 0)
		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("数据库未初始化")
	}

	cacheMutex.Lock()
	defer cacheMutex.Unlock()

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

	return err
}

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

	cacheMutex.Lock()
	defer cacheMutex.Unlock()

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

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

	cacheMutex.Lock()
	defer cacheMutex.Unlock()

	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, key string, ver int) error {
	pool := GetDeviceKeyPool()
	if !pool.IsEnabled() {
		return nil
	}

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

	for _, item := range pool.items {
		if item.DeviceHeader == deviceHeader {
			item.Key = key
			item.KeyVer = ver
			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
}