package com.easyagents.agent.runtime.event.interceptor;

import com.easyagents.agent.runtime.AgentRuntimeException;
import com.easyagents.agent.runtime.event.AgentRuntimeEvent;
import com.easyagents.agent.runtime.event.AgentRuntimeEventBridge;
import com.easyagents.agent.runtime.event.AgentRuntimeEventType;
import com.easyagents.agent.runtime.event.AgentRuntimeInterceptor;
import io.agentscope.core.ReActAgent;
import io.agentscope.core.agent.Agent;
import io.agentscope.core.hook.HookEvent;
import io.agentscope.core.hook.PreCallEvent;
import io.agentscope.core.hook.PreReasoningEvent;
import io.agentscope.core.memory.Memory;
import io.agentscope.core.memory.autocontext.*;
import io.agentscope.core.message.Msg;
import io.agentscope.core.message.MsgRole;
import io.agentscope.core.message.TextBlock;
import io.agentscope.core.plan.PlanNotebook;
import io.agentscope.core.tool.Toolkit;
import reactor.core.publisher.Mono;

import java.util.ArrayList;
import java.util.LinkedHashMap;
import java.util.List;
import java.util.Map;
import java.util.concurrent.atomic.AtomicBoolean;

/**
 * AutoContext 主线路干预器。
 *
 * <p><strong>特殊约束：</strong>本干预器是对 AgentScope 官方
 * {@code AutoContextHook} 的替代实现，不能与官方 {@code AutoContextHook} 同时注册。
 * 两者同时存在会导致 {@link AutoContextMemory#compressIfNeeded()}、上下文重写以及
 * {@link ContextOffloadTool} 注册被重复执行。</p>
 *
 * <p>本类承担两类职责。第一类是主线路干预：在 {@link PreCallEvent} 中注册
 * AutoContext 工具能力，在 {@link PreReasoningEvent} 中触发记忆压缩并改写
 * LLM 输入消息。第二类是旁路通知：通过 {@link AgentRuntimeEventBridge} 发出
 * {@link AgentRuntimeEventType#MEMORY_COMPRESSION_STARTED} 和
 * {@link AgentRuntimeEventType#MEMORY_COMPRESSION_COMPLETED}，这些事件只用于调用方展示，
 * 不写入 AgentScope memory/session，也不参与 LLM 会话协议。</p>
 */
public class AutoContextInterceptor implements AgentRuntimeInterceptor {

    private static final String STATUS_KEY = "memory-compression";
    private static final String AUTO_CONTEXT_SYSTEM_INSTRUCTION =
            "You may see compressed messages containing <!-- CONTEXT_OFFLOAD uuid=... -->.\n"
                    + "- Use the UUID to call context_reload if you need full details.\n"
                    + "- NEVER mention, quote, or refer to UUIDs, offload tags, or internal metadata in your response.";

    private final AgentRuntimeEventBridge eventBridge;
    private final AutoContextConfig autoContextConfig;
    private final AtomicBoolean registered = new AtomicBoolean(false);

    /**
     * 创建 AutoContext 主线路干预器。
     *
     * <p>传入的 {@link AutoContextConfig} 必须是创建目标 {@link AutoContextMemory}
     * 时使用的同一份配置。这样压缩开始事件的触发条件才能与 AgentScope
     * {@code compressIfNeeded()} 的入口判断保持一致。</p>
     *
     * @param eventBridge 旁路事件桥
     * @param autoContextConfig AutoContext 配置
     */
    public AutoContextInterceptor(AgentRuntimeEventBridge eventBridge,
                                  AutoContextConfig autoContextConfig) {
        this.eventBridge = eventBridge;
        this.autoContextConfig = autoContextConfig;
    }

    /**
     * 处理 AutoContext 相关 AgentScope Hook 事件。
     *
     * @param event AgentScope Hook 事件
     * @param <T> Hook 事件类型
     * @return 处理后的 Hook 事件
     */
    @Override
    public <T extends HookEvent> Mono<T> intercept(T event) {
        if (event instanceof PreCallEvent preCallEvent) {
            @SuppressWarnings("unchecked")
            Mono<T> result = (Mono<T>) handlePreCall(preCallEvent);
            return result;
        }
        if (event instanceof PreReasoningEvent preReasoningEvent) {
            @SuppressWarnings("unchecked")
            Mono<T> result = (Mono<T>) handlePreReasoning(preReasoningEvent);
            return result;
        }
        return Mono.just(event);
    }

    /**
     * 获取干预器优先级。
     *
     * @return 优先级，保持与官方 AutoContextHook 一致
     */
    @Override
    public int priority() {
        return 0;
    }

    /**
     * 判断 AutoContext 工具集成是否已经注册。
     *
     * @return 已注册时返回 true
     */
    public boolean isRegistered() {
        return registered.get();
    }

    /**
     * 处理调用前事件，注册 AutoContext 的上下文重载工具和计划本。
     *
     * @param event 调用前事件
     * @return 原事件
     */
    private Mono<PreCallEvent> handlePreCall(PreCallEvent event) {
        if (registered.get()) {
            return Mono.just(event);
        }
        Agent agent = event.getAgent();
        if (!(agent instanceof ReActAgent reActAgent)) {
            return Mono.just(event);
        }
        Memory memory = reActAgent.getMemory();
        if (!(memory instanceof AutoContextMemory autoContextMemory)) {
            return Mono.just(event);
        }
        if (!registered.compareAndSet(false, true)) {
            return Mono.just(event);
        }
        try {
            Toolkit toolkit = reActAgent.getToolkit();
            if (toolkit != null) {
                toolkit.registerTool(new ContextOffloadTool(autoContextMemory));
            }
            PlanNotebook planNotebook = reActAgent.getPlanNotebook();
            if (planNotebook != null) {
                autoContextMemory.attachPlanNote(planNotebook);
            }
        } catch (Exception e) {
            registered.set(false);
            throw new AgentRuntimeException("Failed to register AutoContext integration.", e);
        }
        return Mono.just(event);
    }

    /**
     * 处理推理前事件，触发 AutoContext 压缩并改写 LLM 输入消息。
     *
     * <p>这是主线路干预逻辑：{@code compressIfNeeded()} 和 {@code setInputMessages(...)}
     * 会影响 AgentScope 本次 reasoning 输入。压缩开始/完成事件则是旁路通知，只发给调用方。</p>
     *
     * @param event 推理前事件
     * @return 改写输入消息后的事件
     */
    private Mono<PreReasoningEvent> handlePreReasoning(PreReasoningEvent event) {
        Agent agent = event.getAgent();
        if (!(agent instanceof ReActAgent reActAgent)) {
            return Mono.just(event);
        }
        Memory memory = reActAgent.getMemory();
        if (!(memory instanceof AutoContextMemory autoContextMemory)) {
            return Mono.just(event);
        }
        // 判断是否达到压缩条件
        CompressionCheck compressionCheck = compressionCheck(autoContextMemory.getMessages());
        int beforeEventCount = compressionEventCount(autoContextMemory);
        if (compressionCheck.thresholdReached()) {
            emitCompressionStarted(compressionCheck);
            boolean compressed = autoContextMemory.compressIfNeeded();
            List<CompressionEvent> newEvents = newCompressionEvents(autoContextMemory, beforeEventCount);
            emitCompressionCompleted(compressionCheck, newEvents, compressed && !newEvents.isEmpty());
        }
        // 压缩完毕后自动进行当前的会话
        event.setInputMessages(buildInputMessages(event, autoContextMemory));
        return Mono.just(event);
    }

    /**
     * 计算当前记忆是否达到 AutoContext 官方压缩入口条件。
     *
     * @param messages 当前工作记忆消息
     * @return 压缩入口检查结果
     */
    private CompressionCheck compressionCheck(List<Msg> messages) {
        List<Msg> safeMessages = messages == null ? List.of() : messages;
        int messageCount = safeMessages.size();
        int tokenCount = TokenCounterUtil.calculateToken(safeMessages);
        int thresholdMessageCount = autoContextConfig == null ? Integer.MAX_VALUE : autoContextConfig.getMsgThreshold();
        long thresholdTokenCount = autoContextConfig == null
                ? Long.MAX_VALUE
                : (long) (autoContextConfig.getMaxToken() * autoContextConfig.getTokenRatio());
        boolean messageThresholdReached = messageCount >= thresholdMessageCount;
        boolean tokenThresholdReached = tokenCount >= thresholdTokenCount;
        return new CompressionCheck(messageCount, tokenCount, thresholdMessageCount, thresholdTokenCount,
                messageThresholdReached, tokenThresholdReached);
    }

    /**
     * 获取当前压缩事件数量。
     *
     * @param autoContextMemory AutoContext 记忆
     * @return 压缩事件数量
     */
    private int compressionEventCount(AutoContextMemory autoContextMemory) {
        List<CompressionEvent> events = autoContextMemory.getCompressionEvents();
        return events == null ? 0 : events.size();
    }

    /**
     * 获取本次压缩新增的 AgentScope 压缩事件。
     *
     * @param autoContextMemory AutoContext 记忆
     * @param beforeEventCount 压缩前事件数量
     * @return 新增压缩事件
     */
    private List<CompressionEvent> newCompressionEvents(AutoContextMemory autoContextMemory, int beforeEventCount) {
        List<CompressionEvent> events = autoContextMemory.getCompressionEvents();
        if (events == null || events.size() <= beforeEventCount) {
            return List.of();
        }
        return new ArrayList<>(events.subList(beforeEventCount, events.size()));
    }

    /**
     * 构建 AutoContext 压缩后传给 LLM 的输入消息。
     *
     * @param event 推理前事件
     * @param autoContextMemory AutoContext 记忆
     * @return 更新后的输入消息
     */
    private List<Msg> buildInputMessages(PreReasoningEvent event, AutoContextMemory autoContextMemory) {
        List<Msg> originalInputMessages = event.getInputMessages();
        List<Msg> newInputMessages = new ArrayList<>();
        if (!originalInputMessages.isEmpty() && originalInputMessages.get(0).getRole() == MsgRole.SYSTEM) {
            Msg originalSystemMsg = originalInputMessages.get(0);
            String originalSystemText = originalSystemMsg.getTextContent();
            String newSystemText = originalSystemText != null
                    ? originalSystemText + "\n\n" + AUTO_CONTEXT_SYSTEM_INSTRUCTION
                    : AUTO_CONTEXT_SYSTEM_INSTRUCTION;
            newInputMessages.add(Msg.builder()
                    .role(MsgRole.SYSTEM)
                    .name(originalSystemMsg.getName())
                    .content(TextBlock.builder().text(newSystemText).build())
                    .metadata(originalSystemMsg.getMetadata())
                    .build());
        } else {
            newInputMessages.add(Msg.builder()
                    .role(MsgRole.SYSTEM)
                    .name("system")
                    .content(TextBlock.builder().text(AUTO_CONTEXT_SYSTEM_INSTRUCTION).build())
                    .build());
        }
        newInputMessages.addAll(autoContextMemory.getMessages());
        return newInputMessages;
    }

    /**
     * 发射上下文压缩开始旁路事件。
     *
     * @param compressionCheck 压缩入口检查结果
     */
    private void emitCompressionStarted(CompressionCheck compressionCheck) {
        AgentRuntimeEvent event = eventBridge.event(AgentRuntimeEventType.MEMORY_COMPRESSION_STARTED);
        event.getPayload().put("statusKey", STATUS_KEY);
        event.getPayload().put("phase", "started");
        event.getPayload().put("status", "running");
        event.getPayload().put("label", "正在整理上下文");
        putCompressionCheckPayload(event, compressionCheck);
        eventBridge.emit(event);
    }

    /**
     * 发射上下文压缩完成旁路事件。
     *
     * @param compressionCheck 压缩入口检查结果
     * @param events 新增压缩事件
     */
    private void emitCompressionCompleted(CompressionCheck compressionCheck,
                                          List<CompressionEvent> events,
                                          boolean compressed) {
        AgentRuntimeEvent event = eventBridge.event(AgentRuntimeEventType.MEMORY_COMPRESSION_COMPLETED);
        event.getPayload().put("statusKey", STATUS_KEY);
        event.getPayload().put("phase", "completed");
        event.getPayload().put("status", "done");
        event.getPayload().put("label", compressed ? "已整理上下文" : "无需压缩上下文");
        event.getPayload().put("compressed", compressed);
        event.getPayload().put("eventCount", events == null ? 0 : events.size());
        event.getPayload().put("events", toPayloadEvents(events));
        putCompressionCheckPayload(event, compressionCheck);
        eventBridge.emit(event);
    }

    /**
     * 填充压缩入口判断相关载荷。
     *
     * @param event 运行时事件
     * @param compressionCheck 压缩入口检查结果
     */
    private void putCompressionCheckPayload(AgentRuntimeEvent event, CompressionCheck compressionCheck) {
        event.getPayload().put("messageCount", compressionCheck.messageCount());
        event.getPayload().put("tokenCount", compressionCheck.tokenCount());
        event.getPayload().put("thresholdMessageCount", compressionCheck.thresholdMessageCount());
        event.getPayload().put("thresholdTokenCount", compressionCheck.thresholdTokenCount());
        event.getPayload().put("messageThresholdReached", compressionCheck.messageThresholdReached());
        event.getPayload().put("tokenThresholdReached", compressionCheck.tokenThresholdReached());
    }

    /**
     * 将 AgentScope 压缩事件转换为旁路事件载荷。
     *
     * @param events AgentScope 压缩事件
     * @return 压缩事件载荷
     */
    private List<Map<String, Object>> toPayloadEvents(List<CompressionEvent> events) {
        if (events == null || events.isEmpty()) {
            return List.of();
        }
        List<Map<String, Object>> payloadEvents = new ArrayList<>();
        for (CompressionEvent event : events) {
            Map<String, Object> payload = new LinkedHashMap<>();
            payload.put("eventType", event.getEventType());
            payload.put("timestamp", event.getTimestamp());
            payload.put("compressedMessageCount", event.getCompressedMessageCount());
            payload.put("previousMessageId", event.getPreviousMessageId());
            payload.put("nextMessageId", event.getNextMessageId());
            payload.put("compressedMessageId", event.getCompressedMessageId());
            payload.put("tokenBefore", event.getTokenBefore());
            payload.put("tokenAfter", event.getTokenAfter());
            payload.put("tokenReduction", event.getTokenReduction());
            payload.put("inputToken", event.getCompressInputToken());
            payload.put("outputToken", event.getCompressOutputToken());
            payloadEvents.add(payload);
        }
        return payloadEvents;
    }

    private record CompressionCheck(int messageCount,
                                    int tokenCount,
                                    int thresholdMessageCount,
                                    long thresholdTokenCount,
                                    boolean messageThresholdReached,
                                    boolean tokenThresholdReached) {

        /**
         * 判断是否达到 AutoContext 官方压缩入口条件。
         *
         * @return 达到任一阈值时返回 true
         */
        private boolean thresholdReached() {
            return messageThresholdReached || tokenThresholdReached;
        }
    }
}
