common/redis/redis.go

package redis

import (
	"context"
	"strings"
	"sync"

	"github.com/gogf/gf/v2/database/gredis"
	"github.com/gogf/gf/v2/frame/g"
	"github.com/gogf/gf/v2/os/glog"
	"github.com/gogf/gf/v2/os/gtime"
	"github.com/gogf/gf/v2/util/gconv"
)

var (
	// redisClient 内部使用的 Redis 客户端（单例模式）
	redisClient *gredis.Redis
	redisOnce   sync.Once
)

// getClient 获取 Redis 客户端（延迟初始化）
func getClient() *gredis.Redis {
	redisOnce.Do(func() {
		redisClient = g.Redis()
	})
	return redisClient
}

// GetRedisClient 获取 Redis 客户端（供外部使用）
func GetRedisClient() *gredis.Redis {
	return getClient()
}

// RedisClient 导出的 Redis 客户端（供 mongo.go 使用，兼容旧代码）
var RedisClient = getClient()

// Stream 和消费者组常量
const (
	// RAGFlow 请求 Stream Key
	RAGFlowRequestStreamKey = "ragflow:request:stream"
	// RAGFlow 响应 Stream Key
	RAGFlowResponseStreamKey = "ragflow:response:stream"
	// RAGFlow 请求消费者组名称
	RAGFlowRequestConsumerGroup = "ragflow:request:consumer:group"
	// RAGFlow 响应消费者组名称
	RAGFlowResponseConsumerGroup = "ragflow:response:consumer:group"
	// RAGFlow 消费者组名称（兼容旧代码）
	RAGFlowConsumerGroup = "ragflow:consumer:group"
	// 会话最后活跃时间 Key 前缀
	SessionLastActiveKeyPrefix = "ragflow:session:"
)

// StreamMessage Redis Stream 消息结构
type StreamMessage struct {
	ID     string                 // 消息ID（自动生成）
	Values map[string]interface{} // 消息内容
}

// InitStreamGroup 初始化 Stream 和消费者组
// 使用 gredis Do() 方法执行 XGROUP CREATE 命令
func InitStreamGroup(ctx context.Context, streamKey, groupName string) error {
	// XGROUP CREATE streamKey groupName 0 MKSTREAM
	_, err := redisClient.Do(ctx, "XGROUP", "CREATE", streamKey, groupName, "0", "MKSTREAM")
	if err != nil {
		// 如果组已存在，忽略错误
		errStr := err.Error()
		if strings.Contains(errStr, "BUSYGROUP") || strings.Contains(errStr, "already exists") {
			return nil
		}
		return err
	}
	return nil
}

// AddToStream 将消息添加到 Stream
// 使用 gredis Do() 方法执行 XADD 命令
// msg 可以是结构体或 map，内部自动转换
func AddToStream(ctx context.Context, streamKey string, msg interface{}) (messageID string, err error) {
	// 将结构体转换为 map
	values := gconv.Map(msg)

	// XADD streamKey * field1 value1 field2 value2 ...
	args := make([]interface{}, 0, len(values)*2+2)
	args = append(args, streamKey, "*") // "*" 自动生成ID
	for key, val := range values {
		args = append(args, key, val)
	}

	result, err := redisClient.Do(ctx, "XADD", args...)
	if err != nil {
		return
	}

	messageID = result.String()
	return
}

// CreateConsumerGroup 创建消费者组（如果不存在）
// XGROUP CREATE streamKey groupName $ MKSTREAM
func CreateConsumerGroup(ctx context.Context, streamKey, groupName string) error {
	_, err := redisClient.Do(ctx, "XGROUP", "CREATE", streamKey, groupName, "$", "MKSTREAM")
	return err
}

// ReadFromStream 从 Stream 读取消息（消费者组模式）
// 使用 gredis Do() 方法执行 XREADGROUP 命令
func ReadFromStream(ctx context.Context, streamKey, groupName, consumerName string, count int64, blockMs int64) ([]StreamMessage, error) {
	glog.Debugf(ctx, "[DEBUG Redis] XREADGROUP GROUP %s %s COUNT %d BLOCK %d STREAMS %s >",
		groupName, consumerName, count, blockMs, streamKey)

	// XREADGROUP GROUP groupName consumerName COUNT count BLOCK blockMs STREAMS streamKey >
	result, err := redisClient.Do(ctx,
		"XREADGROUP", "GROUP", groupName, consumerName,
		"COUNT", count,
		"BLOCK", blockMs,
		"STREAMS", streamKey, ">",
	)

	if err != nil {
		glog.Errorf(ctx, "[DEBUG Redis] XREADGROUP 错误: %v", err)
		return nil, err
	}

	glog.Debugf(ctx, "[DEBUG Redis] XREADGROUP 返回: %+v", result)

	// 预分配容量，避免动态扩容
	messages := make([]StreamMessage, 0, int(count))

	if result == nil || result.IsEmpty() {
		// 超时或没有数据
		return messages, nil
	}

	// GoFrame gredis 返回格式: map[streamKey:[[msgID [field1 value1 field2 value2 ...]] ...]]
	resultVal := result.Val()

	// 尝试 map 格式（GoFrame gredis 返回）
	if streamsMap, ok := resultVal.(map[interface{}]interface{}); ok {
		for _, streamMsgs := range streamsMap {
			msgsArray, ok := streamMsgs.([]interface{})
			if !ok {
				continue
			}
			for _, msgData := range msgsArray {
				msgArray, ok := msgData.([]interface{})
				if !ok || len(msgArray) < 2 {
					continue
				}
				msgID := gconv.String(msgArray[0])
				fieldsArray, ok := msgArray[1].([]interface{})
				if !ok {
					continue
				}
				values := make(map[string]interface{}, len(fieldsArray)/2)
				for i := 0; i < len(fieldsArray); i += 2 {
					if i+1 < len(fieldsArray) {
						key := gconv.String(fieldsArray[i])
						values[key] = fieldsArray[i+1]
					}
				}
				messages = append(messages, StreamMessage{
					ID:     msgID,
					Values: values,
				})
			}
		}
		return messages, nil
	}

	// 尝试数组格式（标准 Redis 返回）
	if streamsArray, ok := resultVal.([]interface{}); ok && len(streamsArray) > 0 {
		for _, streamData := range streamsArray {
			streamArray, ok := streamData.([]interface{})
			if !ok || len(streamArray) < 2 {
				continue
			}
			messagesArray, ok := streamArray[1].([]interface{})
			if !ok {
				continue
			}
			for _, msgData := range messagesArray {
				msgArray, ok := msgData.([]interface{})
				if !ok || len(msgArray) < 2 {
					continue
				}
				msgID := gconv.String(msgArray[0])
				fieldsArray, ok := msgArray[1].([]interface{})
				if !ok {
					continue
				}
				values := make(map[string]interface{}, len(fieldsArray)/2)
				for i := 0; i < len(fieldsArray); i += 2 {
					if i+1 < len(fieldsArray) {
						key := gconv.String(fieldsArray[i])
						values[key] = fieldsArray[i+1]
					}
				}
				messages = append(messages, StreamMessage{
					ID:     msgID,
					Values: values,
				})
			}
		}
	}

	return messages, nil
}

// AckMessage 确认消息已处理
// 使用 gredis Do() 方法执行 XACK 命令
func AckMessage(ctx context.Context, streamKey, groupName string, messageIDs ...string) error {
	// XACK streamKey groupName messageID1 messageID2 ...
	// 预分配容量，避免动态扩容
	args := make([]interface{}, 0, len(messageIDs)+2)
	args = append(args, streamKey, groupName)
	for _, id := range messageIDs {
		args = append(args, id)
	}

	_, err := redisClient.Do(ctx, "XACK", args...)
	return err
}

// GetStreamLength 获取 Stream 当前长度
// 使用 gredis Do() 方法执行 XLEN 命令
func GetStreamLength(ctx context.Context, streamKey string) (int64, error) {
	// XLEN streamKey
	result, err := redisClient.Do(ctx, "XLEN", streamKey)
	if err != nil {
		return 0, err
	}

	length := gconv.Int64(result)
	return length, nil
}

// PendingMessage Pending 消息结构
type PendingMessage struct {
	ID         string // 消息ID
	Consumer   string // 消费者名称
	Idle       int64  // 空闲时间(毫秒)
	RetryCount int64  // 重试次数
}

// GetPendingMessages 获取待处理消息
// 使用 gredis Do() 方法执行 XPENDING 命令
func GetPendingMessages(ctx context.Context, streamKey, groupName string, start, end string, count int64) ([]PendingMessage, error) {
	// XPENDING streamKey groupName start end count
	result, err := redisClient.Do(ctx, "XPENDING", streamKey, groupName, start, end, count)
	if err != nil {
		return nil, err
	}

	if result == nil {
		return nil, nil
	}

	// 解析返回值：[[ID, consumer, idle, retryCount], ...]
	pendingArray, ok := result.Val().([]interface{})
	if !ok {
		return nil, nil
	}

	messages := make([]PendingMessage, 0, len(pendingArray))
	for _, item := range pendingArray {
		itemArray, ok := item.([]interface{})
		if !ok || len(itemArray) < 4 {
			continue
		}

		messages = append(messages, PendingMessage{
			ID:         gconv.String(itemArray[0]),
			Consumer:   gconv.String(itemArray[1]),
			Idle:       gconv.Int64(itemArray[2]),
			RetryCount: gconv.Int64(itemArray[3]),
		})
	}

	return messages, nil
}

// ClaimPendingMessage 认领超时的 Pending 消息
// 使用 gredis Do() 方法执行 XCLAIM 命令
func ClaimPendingMessage(ctx context.Context, streamKey, groupName, consumerName string, minIdleTime int64, messageIDs ...string) ([]StreamMessage, error) {
	// XCLAIM streamKey groupName consumerName minIdleTime messageID1 messageID2 ...
	args := []interface{}{streamKey, groupName, consumerName, minIdleTime}
	for _, id := range messageIDs {
		args = append(args, id)
	}

	result, err := redisClient.Do(ctx, "XCLAIM", args...)
	if err != nil {
		return nil, err
	}

	if result == nil {
		return nil, nil
	}

	// 解析返回值：类似 XREADGROUP
	messagesArray, ok := result.Val().([]interface{})
	if !ok {
		return nil, nil
	}

	// 预分配容量,避免动态扩容
	messages := make([]StreamMessage, 0, len(messagesArray))
	for _, msgData := range messagesArray {
		msgArray, ok := msgData.([]interface{})
		if !ok || len(msgArray) < 2 {
			continue
		}

		msgID := gconv.String(msgArray[0])
		fieldsArray, ok := msgArray[1].([]interface{})
		if !ok {
			continue
		}

		// 预分配 map 容量 ,避免动态扩容
		values := make(map[string]interface{}, len(fieldsArray)/2)
		for i := 0; i < len(fieldsArray); i += 2 {
			if i+1 < len(fieldsArray) {
				key := gconv.String(fieldsArray[i])
				values[key] = fieldsArray[i+1]
			}
		}

		messages = append(messages, StreamMessage{
			ID:     msgID,
			Values: values,
		})
	}

	return messages, nil
}

// SetSessionLastActive 设置用户最后活跃时间
// 使用 gredis SetEX 方法
func SetSessionLastActive(ctx context.Context, userId string) error {
	key := SessionLastActiveKeyPrefix + userId + ":last_active"
	timestamp := gtime.Now().Timestamp()

	// SETEX key 7200 value (7200秒 = 2小时)
	_, err := redisClient.Do(ctx, "SETEX", key, 7200, timestamp)
	return err
}

// GetSessionLastActive 获取用户最后活跃时间
// 使用 gredis Get 方法
func GetSessionLastActive(ctx context.Context, userId string) (int64, error) {
	key := SessionLastActiveKeyPrefix + userId + ":last_active"
	result, err := redisClient.Get(ctx, key)
	if err != nil {
		return 0, err
	}

	if result.IsEmpty() {
		return 0, nil
	}

	timestamp := gconv.Int64(result.Val())
	return timestamp, nil
}

// IsUserActive 检查用户是否在指定时间范围内活跃过
// 用于追问逻辑：如果用户最近活跃过，则不发送追问消息
// 参数：
//   - userId: 用户ID
//   - seconds: 时间范围（秒），例如传入300表示检查5分钟内是否活跃
//
// 返回：
//   - bool: true表示用户在指定时间内活跃过
//   - error: 操作失败时返回错误
func IsUserActive(ctx context.Context, userId string, seconds int64) (bool, error) {
	lastActive, err := GetSessionLastActive(ctx, userId)
	if err != nil {
		return false, err
	}

	if lastActive == 0 {
		return false, nil // 未找到记录，视为不活跃
	}

	now := gtime.Now().Timestamp()
	return (now - lastActive) < seconds, nil
}

// SetSessionCache 缓存用户的 RAGFlow Session ID
// 使用 gredis SetEX 方法
func SetSessionCache(ctx context.Context, userId, sessionId string) error {
	key := SessionLastActiveKeyPrefix + userId + ":session_id"

	// SETEX key 604800 value (604800秒 = 7天)
	_, err := redisClient.Do(ctx, "SETEX", key, 604800, sessionId)
	return err
}

// 限流相关常量
const (
	// RateLimitKeyPrefix 限流计数器 Key 前缀
	RateLimitKeyPrefix = "ragflow:ratelimit:"
)

// IncrRateLimit 增加限流计数器，返回当前计数
// windowSeconds: 时间窗口（秒）
func IncrRateLimit(ctx context.Context, key string, windowSeconds int64) (count int64, err error) {
	fullKey := RateLimitKeyPrefix + key
	result, err := redisClient.Do(ctx, "INCR", fullKey)
	if err != nil {
		return
	}
	count = result.Int64()

	// 首次设置过期时间
	if count == 1 {
		redisClient.Do(ctx, "EXPIRE", fullKey, windowSeconds)
	}
	return
}

// GetRateLimit 获取当前限流计数
func GetRateLimit(ctx context.Context, key string) (count int64, err error) {
	fullKey := RateLimitKeyPrefix + key
	result, err := redisClient.Get(ctx, fullKey)
	if err != nil {
		return
	}
	if result.IsEmpty() {
		return 0, nil
	}
	count = result.Int64()
	return
}

// GetSessionCache 获取缓存的 RAGFlow Session ID
func GetSessionCache(ctx context.Context, userId string) (string, error) {
	key := SessionLastActiveKeyPrefix + userId + ":session_id"
	result, err := redisClient.Get(ctx, key)
	if err != nil {
		return "", err
	}

	if result.IsEmpty() {
		return "", nil
	}

	return result.String(), nil
}

// DelSessionCache 删除缓存的 RAGFlow Session ID（归档时调用）
func DelSessionCache(ctx context.Context, userId string) error {
	key := SessionLastActiveKeyPrefix + userId + ":session_id"
	_, err := redisClient.Del(ctx, key)
	return err
}

// TryLock 尝试获取分布式锁（非阻塞）
// key: 锁的键名
// expireSeconds: 锁的过期时间（秒），防止死锁
// 返回 true 表示获取成功，false 表示锁已被其他节点持有
func TryLock(ctx context.Context, key string, expireSeconds int) bool {
	// SET key value NX EX expireSeconds
	result, err := redisClient.Do(ctx, "SET", key, gtime.Now().String(), "NX", "EX", expireSeconds)
	if err != nil {
		glog.Errorf(ctx, "获取分布式锁失败: %v", err)
		return false
	}
	return result.String() == "OK"
}

// Unlock 释放分布式锁
func Unlock(ctx context.Context, key string) {
	if _, err := redisClient.Del(ctx, key); err != nil {
		glog.Errorf(ctx, "释放分布式锁失败: %v", err)
	}
}