package cn.teammodel.ai.cache;

import cn.hutool.cache.CacheUtil;
import cn.hutool.cache.impl.TimedCache;
import cn.hutool.core.collection.ListUtil;
import cn.teammodel.model.entity.ai.ChatSession.Message;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.lang3.ObjectUtils;

import java.util.List;

/**
 * 聊天记录上下文的缓存
 * @author winter
 * @create 2023-12-20 11:02
 */
@Slf4j
@SuppressWarnings("unchecked")
public class HistoryCache {
    private static TimedCache<Object, Object> HISTORY;
    private static Integer contextSize = 3;

    /**
    * 初始化缓存
    */
    public static void init(Long timeout, Integer contextNum) {
        contextSize = contextNum;
        HISTORY  = CacheUtil.newTimedCache(timeout);
        // 一分钟清理一次
        HISTORY.schedulePrune(60 * 1000);
    }

    public static List<Message>  getContext(String sessionId) {
        return (List<Message>) HISTORY.get(sessionId);
    }

    public static void putContext(String sessionId, List<Message> context) {
        HISTORY.put(sessionId, context);
    }

    public static void removeContext(String sessionId) {
        HISTORY.remove(sessionId);}

    /**
    * 更新上下文, 保证上下文的数量在 contextSize 之内
    */
    public static void updateContext(String sessionId, Message message) {
        List<Message> messages = (List<Message>)HISTORY.get(sessionId);

        if (ObjectUtils.isEmpty(messages)) {
            List<Message> context = ListUtil.of(message);
            HISTORY.put(sessionId, context);
        } else if (messages.size() >= contextSize) {
            // 队列
            messages.remove(0);
            messages.add(message);
        } else {
            messages.add(message);
        }
    }

}
