package com.easyagents.agent.runtime.agentscope;

import com.easyagents.agent.runtime.memory.AgentMemoryCompressionParameter;
import com.easyagents.agent.runtime.memory.AgentMemoryPolicy;
import com.easyagents.agent.runtime.memory.AgentMemorySnapshot;
import com.easyagents.agent.runtime.message.AgentMessage;
import io.agentscope.core.memory.InMemoryMemory;
import io.agentscope.core.memory.Memory;
import io.agentscope.core.memory.autocontext.AutoContextConfig;
import io.agentscope.core.memory.autocontext.AutoContextMemory;
import io.agentscope.core.model.Model;

import java.util.List;

/**
 * 将运行时记忆设置适配为 AgentScope 记忆。
 */
public class AgentScopeMemoryAdapter {

    private final AgentScopeMessageAdapter messageAdapter = new AgentScopeMessageAdapter();

    /**
     * 创建记忆并附加恢复后的历史消息。
     *
     * @param snapshot 记忆快照
     * @param policy 记忆策略
     * @param model 用于自动压缩的模型
     * @return AgentScope 记忆
     */
    public Memory createMemory(AgentMemorySnapshot snapshot, AgentMemoryPolicy policy, Model model) {
        return createMemoryResult(snapshot, policy, model).getMemory();
    }

    /**
     * 创建记忆并返回 AutoContext 配置快照。
     *
     * <p>有状态 runtime 需要把同一份 {@link AutoContextConfig} 交给 AutoContext干预器，
     * 用于判断是否进入压缩流程并发出旁路事件。非 AutoContext 记忆不会返回配置。</p>
     *
     * @param snapshot 记忆快照
     * @param policy 记忆策略
     * @param model 用于自动压缩的模型
     * @return 记忆构建结果
     */
    public AgentScopeMemoryBuildResult createMemoryResult(AgentMemorySnapshot snapshot,
                                                          AgentMemoryPolicy policy,
                                                          Model model) {
        AgentMemoryPolicy safePolicy = policy == null ? AgentMemoryPolicy.autoContext() : policy;
        Memory memory;
        AutoContextConfig autoContextConfig = null;
        if (safePolicy.getType() == com.easyagents.agent.runtime.memory.AgentMemoryType.AUTO_CONTEXT
                && safePolicy.getCompressionParameter().isEnabled()) {
            autoContextConfig = toAutoContextConfig(safePolicy.getCompressionParameter());
            memory = new AutoContextMemory(autoContextConfig, new AgentScopeAutoContextCompressionModel(model));
        } else {
            memory = new InMemoryMemory();
        }
        attachMessages(memory, snapshot, safePolicy.getMaxAttachedMessageCount());
        return new AgentScopeMemoryBuildResult(memory, autoContextConfig);
    }

    /**
     * 将自动压缩参数转换为 AgentScope 配置。
     *
     * @param parameter 自动压缩参数
     * @return AgentScope 配置
     */
    public AutoContextConfig toAutoContextConfig(AgentMemoryCompressionParameter parameter) {
        AgentMemoryCompressionParameter safeParameter = parameter == null ? new AgentMemoryCompressionParameter() : parameter;
        Integer compressionThreshold = safeParameter.getMinCompressionTokenThreshold();
        return AutoContextConfig.builder()
                .msgThreshold(safeParameter.getMsgThreshold() == null ? Integer.MAX_VALUE : safeParameter.getMsgThreshold())
                .lastKeep(safeParameter.getLastKeep())
                .tokenRatio(1.0D)
                .maxToken(compressionThreshold == null ? Long.MAX_VALUE : compressionThreshold)
                .largePayloadThreshold(safeParameter.getLargePayloadThreshold())
                .minCompressionTokenThreshold(compressionThreshold == null ? Integer.MAX_VALUE : compressionThreshold)
                .currentRoundCompressionRatio(safeParameter.getCurrentRoundCompressionRatio())
                .minConsecutiveToolMessages(safeParameter.getMinConsecutiveToolMessages())
                .build();
    }

    private void attachMessages(Memory memory, AgentMemorySnapshot snapshot, int maxAttachedMessageCount) {
        if (snapshot == null || snapshot.getMessages().isEmpty()) {
            return;
        }
        List<AgentMessage> messages = snapshot.getMessages();
        int fromIndex = Math.max(0, messages.size() - maxAttachedMessageCount);
        for (AgentMessage message : messages.subList(fromIndex, messages.size())) {
            memory.addMessage(messageAdapter.toMsg(message));
        }
    }
}
