RAG/src/main/java/com/wok/supportbot/cache/SuggestionCache.java

package com.wok.supportbot.cache;

import lombok.extern.slf4j.Slf4j;
import org.springframework.stereotype.Component;

import java.util.List;
import java.util.Optional;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.ConcurrentHashMap;

/**
 * AI 推荐问题（suggest-message-list）的内存缓存。
 * <p>
 * key 为 chatId，value 为最近一次生成的 3 条建议问题。
 * 不持久化到数据库——suggestions 是实时生成、即时消费的短生命周期数据。
 * 每次新的生成会覆盖前一次的建议。
 * <p>
 * 增加 {@code pending} 映射，避免同一 chatId 的并发请求重复触发 LLM 生成。
 */
@Component
@Slf4j
public class SuggestionCache {

    /** 最大缓存条目数，防止无限增长 */
    private static final int MAX_ENTRIES = 1000;

    /** 生成中任务的并发控制：key=chatId，value=正在执行的生成任务 */
    private final ConcurrentHashMap<String, CompletableFuture<List<String>>> pending = new ConcurrentHashMap<>();

    private final ConcurrentHashMap<String, List<String>> cache = new ConcurrentHashMap<>();

    /**
     * 存入建议问题列表。
     *
     * @param chatId      会话 ID
     * @param suggestions 建议问题列表
     */
    public void put(String chatId, List<String> suggestions) {
        if (chatId == null || suggestions == null || suggestions.isEmpty()) {
            return;
        }
        // 超过容量上限时清理一半旧条目（简单 LRU 近似）
        if (cache.size() >= MAX_ENTRIES) {
            int toRemove = MAX_ENTRIES / 2;
            var it = cache.keySet().iterator();
            while (it.hasNext() && toRemove > 0) {
                it.next();
                it.remove();
                toRemove--;
            }
            log.info("SuggestionCache 达到上限 {}，已清理 {} 条旧记录", MAX_ENTRIES, MAX_ENTRIES / 2);
        }
        cache.put(chatId, List.copyOf(suggestions));
        log.debug("SuggestionCache 写入: chatId={}, count={}", chatId, suggestions.size());
    }

    /**
     * 获取建议问题列表。
     *
     * @param chatId 会话 ID
     * @return 建议问题列表，不存在则返回 Optional.empty()
     */
    public Optional<List<String>> get(String chatId) {
        if (chatId == null) {
            return Optional.empty();
        }
        List<String> suggestions = cache.get(chatId);
        return Optional.ofNullable(suggestions);
    }

    /**
     * 注册一个正在生成中的任务。
     *
     * @param chatId  会话 ID
     * @param future  生成任务
     * @return 若已存在相同 chatId 的任务，返回已存在的任务；否则返回传入的任务
     */
    public CompletableFuture<List<String>> putIfAbsent(String chatId, CompletableFuture<List<String>> future) {
        if (chatId == null || future == null) {
            return future;
        }
        CompletableFuture<List<String>> existing = pending.putIfAbsent(chatId, future);
        if (existing != null) {
            return existing;
        }
        // future 完成后写入缓存并从 pending 移除
        future.whenComplete((result, ex) -> {
            pending.remove(chatId, future);
            if (ex != null) {
                log.warn("SuggestionCache 生成任务异常: chatId={}", chatId, ex);
            } else if (result != null && !result.isEmpty()) {
                put(chatId, result);
            }
        });
        return future;
    }

    /**
     * 清除指定会话的建议缓存（会话删除时调用）。
     *
     * @param chatId 会话 ID
     */
    public void evict(String chatId) {
        if (chatId != null) {
            cache.remove(chatId);
            // 同时取消并清理可能正在进行的生成任务
            CompletableFuture<List<String>> future = pending.remove(chatId);
            if (future != null && !future.isDone()) {
                future.cancel(true);
            }
            log.debug("SuggestionCache 清除: chatId={}", chatId);
        }
    }
}