common/redis/redis.go

package redis

import (
	"context"
	"errors"
	"strings"
	"sync"
	"time"

	"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()
}

// Lock 分布式锁
func Lock(ctx context.Context, key string, expireSeconds int64, fn func(ctx context.Context) error) (success bool, err error) {
	limit := 3
LOOP:
	if limit < 0 {
		return false, errors.New("锁重试次数耗尽")
	}
	limit--
	if val, err := RedisClient.Set(ctx, key, true, gredis.SetOption{
		TTLOption: gredis.TTLOption{
			EX: &expireSeconds,
		},
		NX: true,
	}); err != nil {
		return false, err
	} else {
		if val.Bool() {
			defer func(RedisClient *gredis.Redis, ctx context.Context, key string) {
				if _, err = RedisClient.Del(ctx, key); err != nil {
					glog.Errorf(ctx, "RedisClient.Del error: %v", err)
				}
			}(RedisClient, ctx, key)
			if err = fn(ctx); err != nil {
				return false, err
			}
			return true, nil
		} else {
			time.Sleep(time.Second)
			goto LOOP
		}
	}
}

// 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)

	// 检查是否需要记录trace（避免轮询产生大量trace）
	execCtx := ctx
	if !g.Cfg().MustGet(ctx, "jaeger.traceStream", true).Bool() {
		// 不记录trace：使用background context（不继承span）
		execCtx = context.Background()
	}

	// XREADGROUP GROUP groupName consumerName COUNT count BLOCK blockMs STREAMS streamKey >
	result, err := redisClient.Do(execCtx,
		"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
}

// ============== 限流相关 ==============

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
}

// SetSessionCache 缓存 RAGFlow Session ID（租户+用户隔离）
func SetSessionCache(ctx context.Context, tenantId, userId, sessionId string) error {
	key := SessionLastActiveKeyPrefix + tenantId + ":" + userId + ":session_id"
	// SETEX key 7200 value (7200秒 = 2小时，与last_active保持一致)
	_, err := redisClient.Do(ctx, "SETEX", key, 7200, sessionId)
	return err
}

// GetSessionCache 获取缓存的 RAGFlow Session ID（租户+用户隔离）
func GetSessionCache(ctx context.Context, tenantId, userId string) (string, error) {
	key := SessionLastActiveKeyPrefix + tenantId + ":" + 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, tenantId, userId string) error {
	key := SessionLastActiveKeyPrefix + tenantId + ":" + 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)
	}
}

// ============== 对话计数相关（用于卡片触发）==============

const (
	// UserStateKeyPrefix 用户会话状态 Key 前缀（融合阶段+计数）
	UserStateKeyPrefix = "ragflow:user:state:"
	// UserStateExpireSeconds 用户状态过期时间（5分钟）
	UserStateExpireSeconds = 300
)

// UserState 用户会话状态（阶段+对话计数+咨询方向,统一5分钟过期）
type UserState struct {
	Stage       int    `json:"stage"`       // 当前阶段
	Direction   string `json:"direction"`   // 咨询方向
	Count       int64  `json:"count"`       // 对话计数（v5.2卡片触发）
	AccountName string `json:"accountName"` // 用户选择的方向对应的客服账号名称
}

// GetUserState 获取用户状态（阶段+计数）
func GetUserState(ctx context.Context, userId, platform string) (state *UserState, err error) {
	key := UserStateKeyPrefix + userId + "_" + platform
	result, err := redisClient.Do(ctx, "HGETALL", key)
	if err != nil {
		return
	}

	state = &UserState{Stage: 5} // 默认状态5（未选择方向）
	if result.IsEmpty() {
		// Redis为空，初始化默认状态
		if initErr := SetUserStage(ctx, userId, platform, 5); initErr != nil {
			err = initErr
			return
		}
		return
	}

	m := result.Map()
	state.Stage = gconv.Int(m["stage"])
	state.Count = gconv.Int64(m["count"])
	state.Direction = gconv.String(m["direction"])
	return
}

// SetUserStage 设置用户阶段，并刷新过期时间
func SetUserStage(ctx context.Context, userId, platform string, stage int) error {
	key := UserStateKeyPrefix + userId + "_" + platform
	_, err := redisClient.Do(ctx, "HSET", key, "stage", stage)
	if err != nil {
		return err
	}
	_, err = redisClient.Do(ctx, "EXPIRE", key, UserStateExpireSeconds)
	return err
}

// SetUserAccountName 设置用户对应的客服账号名称，并刷新过期时间
func SetUserAccountName(ctx context.Context, userId, platform, accountName string) error {
	key := UserStateKeyPrefix + userId + "_" + platform
	_, err := redisClient.Do(ctx, "HSET", key, "accountName", accountName)
	if err != nil {
		return err
	}
	_, err = redisClient.Do(ctx, "EXPIRE", key, UserStateExpireSeconds)
	return err
}

// SetUserDirection 设置用户选择的咨询方向，并刷新过期时间
func SetUserDirection(ctx context.Context, userId, platform, direction string) error {
	key := UserStateKeyPrefix + userId + "_" + platform
	_, err := redisClient.Do(ctx, "HSET", key, "direction", direction)
	if err != nil {
		return err
	}
	_, err = redisClient.Do(ctx, "EXPIRE", key, UserStateExpireSeconds)
	return err
}

// IncrUserCount 增加用户对话计数，返回当前轮数，并刷新过期时间
func IncrUserCount(ctx context.Context, userId, platform string) (count int64, err error) {
	key := UserStateKeyPrefix + userId + "_" + platform
	result, err := redisClient.Do(ctx, "HINCRBY", key, "count", 1)
	if err != nil {
		return
	}
	count = result.Int64()
	_, err = redisClient.Do(ctx, "EXPIRE", key, UserStateExpireSeconds)
	return
}

// ResetUserState 重置用户状态（归档时调用）
func ResetUserState(ctx context.Context, userId, platform string) error {
	key := UserStateKeyPrefix + userId + "_" + platform
	_, err := redisClient.Del(ctx, key)
	return err
}

// ============== 对话缓存相关（5句落库）==============

const (
	// ConversationCacheKeyPrefix 对话缓存 Key 前缀
	ConversationCacheKeyPrefix = "ragflow:conversation:cache:"
	// ConversationCacheExpireSeconds 对话缓存过期时间（10分钟）
	ConversationCacheExpireSeconds = 600
)

// CacheConversation 缓存单条对话到Redis List（按sessionId存储）
func CacheConversation(ctx context.Context, sessionId string, data []byte) error {
	key := ConversationCacheKeyPrefix + sessionId
	_, err := redisClient.Do(ctx, "RPUSH", key, string(data))
	if err != nil {
		return err
	}
	_, err = redisClient.Do(ctx, "EXPIRE", key, ConversationCacheExpireSeconds)
	return err
}

// GetCachedConversations 获取缓存的对话列表并清空（按sessionId查询）
func GetCachedConversations(ctx context.Context, sessionId string) (list []string, err error) {
	key := ConversationCacheKeyPrefix + sessionId
	result, err := redisClient.Do(ctx, "LRANGE", key, 0, -1)
	if err != nil {
		return
	}
	if result.IsEmpty() {
		return
	}
	list = result.Strings()
	// 清空缓存
	redisClient.Del(ctx, key)
	return
}

// GetCachedConversationCount 获取缓存的对话数量（按sessionId查询）
func GetCachedConversationCount(ctx context.Context, sessionId string) (count int64, err error) {
	key := ConversationCacheKeyPrefix + sessionId
	result, err := redisClient.Do(ctx, "LLEN", key)
	if err != nil {
		return
	}
	return result.Int64(), nil
}

// ClearCachedConversations 清空对话缓存（归档时调用，按sessionId）
func ClearCachedConversations(ctx context.Context, sessionId string) error {
	key := ConversationCacheKeyPrefix + sessionId
	_, err := redisClient.Del(ctx, key)
	return err
}

// ========== 以下为兼容旧接口（内部调用新实现）==========

// IncrConversationCount 增加用户对话计数（兼容旧接口）
func IncrConversationCount(ctx context.Context, userId, platform string, _ int64) (count int64, err error) {
	return IncrUserCount(ctx, userId, platform)
}

// GetConversationCount 获取用户当前对话轮数（兼容旧接口）
func GetConversationCount(ctx context.Context, userId, platform string) (count int64, err error) {
	state, err := GetUserState(ctx, userId, platform)
	if err != nil {
		return
	}
	return state.Count, nil
}

// ResetConversationCount 重置用户对话计数（兼容旧接口）
func ResetConversationCount(ctx context.Context, userId, platform string) error {
	return ResetUserState(ctx, userId, platform)
}