package com.xly.config;


import dev.langchain4j.memory.ChatMemory;
import dev.langchain4j.memory.chat.ChatMemoryProvider;
import dev.langchain4j.memory.chat.MessageWindowChatMemory;
import dev.langchain4j.data.message.ChatMessage;
import org.springframework.stereotype.Component;

import java.util.List;
import java.util.Map;
import java.util.Objects;
import java.util.concurrent.ConcurrentHashMap;

/**
 * 可操作的ChatMemoryProvider：获取消息对象+清除记忆（指定/全量）+删除单条消息
 * 实现框架原生接口，无缝对接AiServices，线程安全适配生产环境
 */
@Component
public class OperableChatMemoryProvider implements ChatMemoryProvider {
    // 核心缓存：memoryId -> ChatMemory，保证一个会话/用户对应唯一记忆实例
    private final Map<Object, ChatMemory> memoryCache = new ConcurrentHashMap<>();
    // 记忆最大消息数，根据业务需求调整（原配置为10）
    private static final int MAX_MESSAGE_SIZE = 100;

    /**
     * 框架原生方法：获取【当前memoryId对应的ChatMemory实例（含消息对象）】
     * AiServices自动调用，也是手动操作记忆/消息的唯一入口
     */
    @Override
    public ChatMemory get(Object memoryId) {
        // 空memoryId兜底，避免空指针
        Object finalMemId = Objects.isNull(memoryId) ? "default_erp_chat_memory" : memoryId;
        // 不存在则创建MessageWindowChatMemory，存在则复用
        return memoryCache.computeIfAbsent(finalMemId, k -> MessageWindowChatMemory.withMaxMessages(MAX_MESSAGE_SIZE));
    }

    // ===================== 1. 获取消息对象（当前memoryId） =====================
    /**
     * 获取当前会话/用户的全部消息列表（含用户消息、AI消息，按对话顺序）
     */
    public List<ChatMessage> getCurrentChatMessages(Object memoryId) {
        if (Objects.isNull(memoryId)) {
            return List.of();
        }
        // 从ChatMemory中获取原生消息列表
        return this.get(memoryId).messages();
    }

    // ===================== 2. 清除记忆（核心需求） =====================
    /**
     * 清空【指定memoryId】的全部记忆（最常用，如前端「清空对话」）
     */
    public void clearSpecifiedMemory(Object memoryId) {
        if (Objects.nonNull(memoryId)) {
            // 调用ChatMemory原生clear()，清空该实例所有消息
            this.get(memoryId).clear();
        }
    }

    /**
     * 全量清除【所有memoryId】的记忆（系统级清理，如定时任务/后台操作）
     */
    public void clearAllMemory() {
        // 遍历所有记忆实例，逐个调用原生clear()清空
        memoryCache.values().forEach(ChatMemory::clear);
        // 可选：彻底清空缓存，销毁所有实例（释放内存，后续会重新创建）
        // memoryCache.clear();
    }

    // ===================== 3. 精准操作：删除单条消息 =====================
    /**
     * 删除当前memoryId的指定单条消息（如删除错误消息）
     * @param message 要删除的消息对象（从getCurrentChatMessages中获取）
     */
    public void deleteSingleMessage(Object memoryId, ChatMessage message) {
        if (Objects.nonNull(memoryId) && Objects.nonNull(message)) {
            ChatMemory currentMemory = this.get(memoryId);
            // 删除指定消息
            currentMemory.messages().remove(message);
            // 刷新记忆，保证MessageWindowChatMemory的最大消息数限制生效
//            currentMemory.update(currentMemory.messages());
        }
    }

    /**
     * 移除并清除指定记忆（清空消息+从缓存删除实例，彻底释放资源，适用于过期会话）
     */
    public void removeAndClearMemory(Object memoryId) {
        if (Objects.nonNull(memoryId)) {
            ChatMemory chatMemory = memoryCache.remove(memoryId);
            if (Objects.nonNull(chatMemory)) {
                chatMemory.clear();
            }
        }
    }
}