package tech.easyflow.agent.runtime.output;

import com.easyagents.agent.runtime.event.AgentRuntimeEvent;
import com.easyagents.agent.runtime.event.AgentRuntimeEventType;
import com.easyagents.agui.AguiExtendedEvent;
import com.easyagents.agui.AguiProtocolEventEncoder;
import com.easyagents.agui.AguiRuntimeEventProjector;
import io.agentscope.core.agui.event.AguiEvent;
import io.agentscope.core.agui.model.AguiMessage;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;
import tech.easyflow.agent.runtime.agui.ResumableAguiSseEmitter;
import tech.easyflow.agent.runtime.hitl.ToolApprovalInputProjection;
import tech.easyflow.core.chat.protocol.ChatDomain;
import tech.easyflow.core.chat.protocol.ChatType;
import tech.easyflow.core.chat.protocol.sse.ChatSseEmitter;

import java.time.Instant;
import java.util.ArrayList;
import java.util.LinkedHashMap;
import java.util.List;
import java.util.Map;
import java.util.UUID;

/**
 * 将同一 Agent 业务运行投影为原生 AG-UI SSE 的输出实现。
 *
 * <p>所有公开发送方法串行化，避免运行线程与 HITL 恢复线程交错破坏事件顺序。</p>
 */
public final class AguiAgentRunOutput implements AgentRunOutput {

    private static final String ASSISTANT_ROLE = "assistant";
    private static final String REASONING_ROLE = "reasoning";

    private final String threadId;
    private final String runId;
    private final String clientUserMessageId;
    private final String clientUserMessageContent;
    private final ChatSseEmitter delegate;
    private final boolean cancelRunOnDisconnect;
    private final AguiRuntimeEventProjector projector;
    private final AguiProtocolEventEncoder encoder = new AguiProtocolEventEncoder();

    private long customSequence;
    private long assistantMessageSequence;
    private long reasoningMessageSequence;
    private String assistantMessageId;
    private String lastAssistantMessageId;
    private String reasoningMessageId;
    private final StringBuilder assistantText = new StringBuilder();
    private final Map<String, Map<String, Object>> activeSkillInvocations = new LinkedHashMap<>();
    private AgentRuntimeEvent pendingCompletedEvent;

    /**
     * 创建 AG-UI 输出。
     *
     * @param threadId 客户端 thread ID
     * @param runId 客户端 run ID
     * @param clientUserMessageId 本轮客户端用户消息 ID
     */
    public AguiAgentRunOutput(String threadId, String runId, String clientUserMessageId) {
        this(threadId, runId, clientUserMessageId, null, new ChatSseEmitter(), true);
    }

    /**
     * 创建包含本轮用户消息快照信息的 AG-UI 输出。
     *
     * @param threadId 客户端 thread ID
     * @param runId 客户端 run ID
     * @param clientUserMessageId 本轮客户端用户消息 ID
     * @param clientUserMessageContent 本轮客户端用户消息正文
     */
    public AguiAgentRunOutput(
            String threadId,
            String runId,
            String clientUserMessageId,
            String clientUserMessageContent) {
        this(threadId, runId, clientUserMessageId, clientUserMessageContent, new ChatSseEmitter(), true);
    }

    /**
     * 使用指定 SSE 发射器创建 AG-UI 输出，供受控装配和测试使用。
     *
     * @param threadId 客户端 thread ID
     * @param runId 客户端 run ID
     * @param clientUserMessageId 本轮客户端用户消息 ID
     * @param delegate SSE 发射器
     */
    public AguiAgentRunOutput(
            String threadId,
            String runId,
            String clientUserMessageId,
            ChatSseEmitter delegate) {
        this(threadId, runId, clientUserMessageId, null, delegate, true);
    }

    /**
     * 使用指定 SSE 发射器和用户消息快照信息创建 AG-UI 输出。
     *
     * @param threadId 客户端 thread ID
     * @param runId 客户端 run ID
     * @param clientUserMessageId 本轮客户端用户消息 ID
     * @param clientUserMessageContent 本轮客户端用户消息正文
     * @param delegate SSE 发射器
     */
    public AguiAgentRunOutput(
            String threadId,
            String runId,
            String clientUserMessageId,
            String clientUserMessageContent,
            ChatSseEmitter delegate) {
        this(threadId, runId, clientUserMessageId, clientUserMessageContent, delegate, true);
    }

    /**
     * 使用指定连接生命周期策略创建 AG-UI 输出。
     *
     * @param threadId 客户端 thread ID
     * @param runId 客户端 run ID
     * @param clientUserMessageId 本轮客户端用户消息 ID
     * @param clientUserMessageContent 本轮客户端用户消息正文
     * @param delegate SSE 发射器
     * @param cancelRunOnDisconnect 浏览器连接断开时是否取消运行
     */
    public AguiAgentRunOutput(
            String threadId,
            String runId,
            String clientUserMessageId,
            String clientUserMessageContent,
            ChatSseEmitter delegate,
            boolean cancelRunOnDisconnect) {
        this.threadId = requireText(threadId, "threadId");
        this.runId = requireText(runId, "runId");
        this.clientUserMessageId = clientUserMessageId;
        this.clientUserMessageContent = clientUserMessageContent;
        this.delegate = java.util.Objects.requireNonNull(delegate, "delegate cannot be null");
        this.cancelRunOnDisconnect = cancelRunOnDisconnect;
        this.projector = new AguiRuntimeEventProjector(threadId, runId);
    }

    @Override
    public SseEmitter emitter() {
        return delegate.getEmitter();
    }

    /**
     * 获取公开 AG-UI 运行 ID。
     *
     * @return 运行 ID
     */
    public String runId() {
        return runId;
    }

    @Override
    public synchronized boolean emitRuntimeEvent(AgentRuntimeEvent event) {
        if (event == null || event.getEventType() == null || delegate.isClosed()) {
            return !delegate.isClosed();
        }
        AgentRuntimeEventType type = event.getEventType();
        if (type == AgentRuntimeEventType.MESSAGE_DELTA
                || type == AgentRuntimeEventType.REASONING_STARTED
                || type == AgentRuntimeEventType.REASONING_DELTA
                || type == AgentRuntimeEventType.REASONING_COMPLETED) {
            // EasyFlow 的跨 delta <think> 归一化结果通过 emitViewEvent 输出，避免重复和标签泄漏。
            return true;
        }
        if (type == AgentRuntimeEventType.TOOL_APPROVAL_REQUIRED) {
            return emitToolApproval(event);
        }
        if (type == AgentRuntimeEventType.SKILL_CALL
                || type == AgentRuntimeEventType.SKILL_RESULT
                || type == AgentRuntimeEventType.SKILL_FAILED) {
            return emitSkillInvocation(event);
        }
        if (type == AgentRuntimeEventType.SKILL_STEP) {
            return true;
        }
        if (isAsyncToolEvent(type)) {
            return true;
        }
        if (type == AgentRuntimeEventType.KNOWLEDGE_RETRIEVAL) {
            return true;
        }
        if (type == AgentRuntimeEventType.MEMORY_COMPRESSION_STARTED
                || type == AgentRuntimeEventType.MEMORY_COMPRESSION_COMPLETED) {
            return true;
        }
        if (type == AgentRuntimeEventType.SUSPENDED) {
            return true;
        }
        if (type == AgentRuntimeEventType.COMPLETED) {
            pendingCompletedEvent = event;
            return true;
        }
        if (type == AgentRuntimeEventType.CANCELLED) {
            if (!closeActiveSkillInvocations("CANCELLED")) {
                return false;
            }
        } else if (type == AgentRuntimeEventType.FAILED) {
            if (!closeActiveSkillInvocations("INCOMPLETE")) {
                return false;
            }
        }
        if (type == AgentRuntimeEventType.TOOL_CALL
                || type == AgentRuntimeEventType.FAILED
                || type == AgentRuntimeEventType.CANCELLED) {
            if (!closeOpenMessages()) {
                return false;
            }
        }
        for (Object protocolEvent : projector.project(event)) {
            if (!send(protocolEvent)) {
                return false;
            }
        }
        if (type == AgentRuntimeEventType.TOOL_CALL) {
            return emitToolMetadata(event);
        }
        return true;
    }

    @Override
    public synchronized boolean emitViewEvent(ChatDomain domain, ChatType type, Object payload) {
        if (delegate.isClosed()) {
            return false;
        }
        if (type == ChatType.MESSAGE) {
            return emitAssistantDelta(stringValue(payload, "delta"));
        }
        if (type == ChatType.THINKING) {
            return emitReasoningDelta(firstText(
                    stringValue(payload, "delta"), stringValue(payload, "reasoning")));
        }
        if (type == ChatType.INPUT_ACCEPTED) {
            Map<String, Object> value = copyMap(payload);
            value.put("clientMessageId", clientUserMessageId);
            value.put("serverMessageId", value.get("messageId"));
            return emitCustom("easyflow.input.accepted", value, clientUserMessageId);
        }
        if (type == ChatType.CITATIONS) {
            return emitCustom("easyflow.knowledge.citations", copyMap(payload), lastAssistantMessageId);
        }
        if ((type == ChatType.TOOL_CALL || type == ChatType.TOOL_RESULT)
                && Boolean.TRUE.equals(copyMap(payload).get("asyncTool"))) {
            return emitCustom("easyflow.async_tool.status",
                    selectPayload(payload,
                            "asyncTool", "asyncToolName", "input", "label", "name", "output",
                            "phase", "result", "sourceToolCallId", "status", "statusKey", "summary",
                            "taskId", "text", "toolCallId", "toolDisplayName", "toolInput", "toolName"),
                    null);
        }
        if (type == ChatType.STATUS) {
            String statusKey = stringValue(payload, "statusKey");
            if ("artifact-published".equals(statusKey)) {
                return emitCustom("easyflow.artifact.published",
                        selectPayload(payload, "schemaVersion", "artifactId", "fileName", "mimeType",
                                "size", "sha256", "downloadUrl", "status"),
                        lastAssistantMessageId);
            }
            if ("knowledge-retrieval".equals(statusKey)) {
                return emitCustom("easyflow.knowledge.retrieval_status",
                        selectPayload(payload, "label", "status", "statusKey"), null);
            }
            if ("memory-compression".equals(statusKey)) {
                return emitCustom("easyflow.runtime.context_status",
                        selectPayload(payload, "compressed", "label", "phase", "status", "statusKey"), null);
            }
        }
        if (type == ChatType.FORM_CANCEL) {
            return emitCustom("easyflow.hitl.tool_approval_resolved", copyMap(payload), null);
        }
        if (type == ChatType.ERROR) {
            return fail(firstText(stringValue(payload, "message"), "Agent runtime failed."), "AGENT_RUNTIME_FAILED");
        }
        // 工具标准事件和业务状态已从 AgentRuntimeEvent 投影；其余 Legacy 展示事件不进入 AG-UI wire。
        return true;
    }

    @Override
    public synchronized boolean canFinishSuccessfully() {
        return pendingCompletedEvent != null;
    }

    @Override
    public boolean cancelRunOnDisconnect() {
        return cancelRunOnDisconnect;
    }

    @Override
    public synchronized boolean finish(String finalText) {
        if (delegate.isClosed()) {
            return false;
        }
        if (!reconcileFinalText(finalText)) {
            return false;
        }
        if (pendingCompletedEvent != null && !projector.isTerminated()) {
            if (!closeActiveSkillInvocations("INCOMPLETE")) {
                return false;
            }
            if (!closeOpenMessages()) {
                return false;
            }
            for (Object protocolEvent : projector.project(pendingCompletedEvent)) {
                if (!send(protocolEvent)) {
                    return false;
                }
            }
            pendingCompletedEvent = null;
        }
        if (!projector.isTerminated()) {
            if (!fail("Agent stream ended without a terminal event.", "MISSING_TERMINAL_EVENT")) {
                return false;
            }
        }
        delegate.complete();
        if (delegate instanceof ResumableAguiSseEmitter resumableEmitter) {
            return resumableEmitter.isJournalCompletionSuccessful();
        }
        return true;
    }

    @Override
    public synchronized void complete() {
        delegate.complete();
    }

    @Override
    public synchronized void completeWithError(Throwable error) {
        if (!projector.isTerminated() && !delegate.isClosed()) {
            closeActiveSkillInvocations("INCOMPLETE");
            fail(error == null || error.getMessage() == null
                    ? "Agent runtime failed."
                    : error.getMessage(), "AGENT_RUNTIME_FAILED");
        }
        delegate.complete();
    }

    @Override
    public boolean isClosed() {
        return delegate.isClosed();
    }

    private boolean emitAssistantDelta(String delta) {
        if (delta == null || delta.isEmpty()) {
            return true;
        }
        if (!ensureRunStarted()) {
            return false;
        }
        if (reasoningMessageId != null && !closeReasoning()) {
            return false;
        }
        if (assistantMessageId == null) {
            assistantMessageId = runId + "-assistant-" + (++assistantMessageSequence);
            lastAssistantMessageId = assistantMessageId;
            if (!send(new AguiEvent.TextMessageStart(
                    threadId, runId, assistantMessageId, ASSISTANT_ROLE))) {
                return false;
            }
        }
        assistantText.append(delta);
        return send(new AguiEvent.TextMessageContent(threadId, runId, assistantMessageId, delta));
    }

    private boolean emitReasoningDelta(String delta) {
        if (delta == null || delta.isEmpty()) {
            return true;
        }
        if (!ensureRunStarted()) {
            return false;
        }
        if (assistantMessageId != null && !closeAssistant()) {
            return false;
        }
        if (reasoningMessageId == null) {
            reasoningMessageId = runId + "-reasoning-" + (++reasoningMessageSequence);
            if (!send(new AguiEvent.ReasoningMessageStart(
                    threadId, runId, reasoningMessageId, REASONING_ROLE))) {
                return false;
            }
        }
        return send(new AguiEvent.ReasoningMessageContent(threadId, runId, reasoningMessageId, delta));
    }

    private boolean emitToolApproval(AgentRuntimeEvent event) {
        Map<String, Object> source = event.getPayload() == null ? Map.of() : event.getPayload();
        Map<String, Object> value = new LinkedHashMap<>();
        value.put("approvalId", event.getMetadata().get("approvalId"));
        value.put("toolCallId", firstText(event.getToolCallId(), stringValue(source, "toolCallId")));
        value.put("toolName", stringValue(source, "toolName"));
        value.put("toolDisplayName", firstText(
                stringValue(source, "toolDisplayName"), stringValue(source, "toolName")));
        value.put("input", ToolApprovalInputProjection.project(
                firstNonNull(source.get("toolInput"), source.get("input"))));
        value.put("expiresAt", source.get("expiresAt"));
        return emitCustom("easyflow.hitl.tool_approval_required", value, event.getMessageId());
    }

    private boolean emitSkillInvocation(AgentRuntimeEvent event) {
        Map<String, Object> value = selectPayload(event.getPayload(),
                "statusKey", "status", "skillId", "skillName", "skillDisplayName",
                "toolCallId", "message");
        String statusKey = stringValue(value, "statusKey");
        if (statusKey == null) {
            return true;
        }
        String status = stringValue(value, "status");
        if ("RUNNING".equals(status)) {
            activeSkillInvocations.put(statusKey, new LinkedHashMap<>(value));
        } else {
            activeSkillInvocations.remove(statusKey);
        }
        return emitCustom("easyflow.skill.invocation_status", value, event.getMessageId());
    }

    private boolean closeActiveSkillInvocations(String status) {
        if (activeSkillInvocations.isEmpty()) {
            return true;
        }
        List<Map<String, Object>> pending = new ArrayList<>(activeSkillInvocations.values());
        activeSkillInvocations.clear();
        for (Map<String, Object> value : pending) {
            Map<String, Object> terminal = new LinkedHashMap<>(value);
            terminal.put("status", status);
            terminal.remove("message");
            if (!emitCustom("easyflow.skill.invocation_status", terminal, null)) {
                return false;
            }
        }
        return true;
    }

    private boolean emitToolMetadata(AgentRuntimeEvent event) {
        Map<String, Object> source = event.getPayload() == null ? Map.of() : event.getPayload();
        String toolCallId = firstText(event.getToolCallId(), stringValue(source, "toolCallId"));
        String toolName = firstText(stringValue(source, "toolName"), stringValue(source, "name"));
        String toolDisplayName = stringValue(source, "toolDisplayName");
        if (toolCallId == null || toolDisplayName == null
                || toolDisplayName.equals(toolName) || isHiddenToolName(toolName)) {
            return true;
        }
        Map<String, Object> value = new LinkedHashMap<>();
        value.put("toolCallId", toolCallId);
        value.put("toolName", toolName);
        value.put("toolDisplayName", toolDisplayName);
        return emitCustom("easyflow.tool.metadata", value, event.getMessageId());
    }

    private boolean emitCustom(String name, Map<String, Object> payload, String messageId) {
        if (!ensureRunStarted()) {
            return false;
        }
        Map<String, Object> value = new LinkedHashMap<>();
        if (payload != null) {
            value.putAll(payload);
        }
        // 协议保留字段由服务端最终写入，避免业务 payload 覆盖运行边界信息。
        value.put("schemaVersion", 1);
        value.put("id", "evt_" + UUID.randomUUID());
        value.put("runId", runId);
        value.put("threadId", threadId);
        value.put("messageId", messageId);
        value.put("sequence", ++customSequence);
        value.put("timestamp", Instant.now().toString());
        return send(new AguiEvent.Custom(threadId, runId, name, value));
    }

    private boolean fail(String message, String code) {
        if (projector.isTerminated()) {
            return true;
        }
        if (!closeOpenMessages()) {
            return false;
        }
        AgentRuntimeEvent failed = AgentRuntimeEvent.of(AgentRuntimeEventType.FAILED);
        failed.getPayload().put("message", message);
        String resolvedCode = code == null || code.isBlank() ? "AGENT_RUNTIME_FAILED" : code;
        for (Object protocolEvent : projector.project(failed)) {
            Object output = protocolEvent instanceof AguiExtendedEvent.RunError runError
                    ? new AguiExtendedEvent.RunError(
                            runError.threadId(), runError.runId(), runError.message(), resolvedCode)
                    : protocolEvent;
            if (!send(output)) {
                return false;
            }
        }
        return true;
    }

    private boolean closeOpenMessages() {
        return closeReasoning() && closeAssistant();
    }

    private boolean closeReasoning() {
        if (reasoningMessageId == null) {
            return true;
        }
        String messageId = reasoningMessageId;
        reasoningMessageId = null;
        return send(new AguiEvent.ReasoningMessageEnd(threadId, runId, messageId));
    }

    private boolean closeAssistant() {
        if (assistantMessageId == null) {
            return true;
        }
        String messageId = assistantMessageId;
        assistantMessageId = null;
        return send(new AguiEvent.TextMessageEnd(threadId, runId, messageId));
    }

    private boolean reconcileFinalText(String finalText) {
        if (finalText == null || finalText.contentEquals(assistantText)) {
            return true;
        }
        String streamedText = assistantText.toString();
        if (finalText.startsWith(streamedText)) {
            return emitAssistantDelta(finalText.substring(streamedText.length()));
        }
        if (!closeOpenMessages()) {
            return false;
        }
        String resolvedAssistantMessageId = lastAssistantMessageId;
        if (resolvedAssistantMessageId == null) {
            resolvedAssistantMessageId = runId + "-assistant-" + (++assistantMessageSequence);
            lastAssistantMessageId = resolvedAssistantMessageId;
        }
        List<AguiMessage> messages = new ArrayList<>(2);
        if (clientUserMessageId != null && clientUserMessageContent != null) {
            messages.add(AguiMessage.userMessage(clientUserMessageId, clientUserMessageContent));
        }
        messages.add(AguiMessage.assistantMessage(resolvedAssistantMessageId, finalText));
        assistantText.setLength(0);
        assistantText.append(finalText);
        return send(new AguiExtendedEvent.MessagesSnapshot(threadId, runId, messages));
    }

    private boolean send(Object event) {
        return delegate.sendData(encoder.encodeToJson(event));
    }

    private boolean ensureRunStarted() {
        if (projector.isTerminated()) {
            return false;
        }
        AgentRuntimeEvent started = AgentRuntimeEvent.of(AgentRuntimeEventType.STARTED);
        for (Object protocolEvent : projector.project(started)) {
            if (!send(protocolEvent)) {
                return false;
            }
        }
        return true;
    }

    private static boolean isAsyncToolEvent(AgentRuntimeEventType type) {
        return type == AgentRuntimeEventType.ASYNC_TOOL_SUBMITTED
                || type == AgentRuntimeEventType.ASYNC_TOOL_OBSERVED
                || type == AgentRuntimeEventType.ASYNC_TOOL_RESULT
                || type == AgentRuntimeEventType.ASYNC_TOOL_CANCELLED
                || type == AgentRuntimeEventType.ASYNC_TOOL_LISTED
                || type == AgentRuntimeEventType.ASYNC_TOOL_FAILED;
    }

    private static boolean isHiddenToolName(String toolName) {
        return "retrieve_knowledge".equalsIgnoreCase(toolName)
                || "context_reload".equalsIgnoreCase(toolName)
                || "__fragment__".equalsIgnoreCase(toolName);
    }

    @SuppressWarnings("unchecked")
    private static Map<String, Object> copyMap(Object payload) {
        return payload instanceof Map<?, ?> map
                ? new LinkedHashMap<>((Map<String, Object>) map)
                : new LinkedHashMap<>();
    }

    /**
     * 选取允许进入 AG-UI CUSTOM 的公开字段。
     *
     * @param payload 服务层展示载荷
     * @param allowedKeys 允许字段
     * @return 公开载荷
     */
    private static Map<String, Object> selectPayload(Object payload, String... allowedKeys) {
        Map<String, Object> source = copyMap(payload);
        Map<String, Object> selected = new LinkedHashMap<>();
        for (String key : allowedKeys) {
            if (source.containsKey(key)) {
                selected.put(key, source.get(key));
            }
        }
        return selected;
    }

    private static String stringValue(Object payload, String key) {
        if (!(payload instanceof Map<?, ?> map)) {
            return null;
        }
        Object value = map.get(key);
        return value instanceof String text ? text : null;
    }

    private static Object firstNonNull(Object first, Object second) {
        return first == null ? second : first;
    }

    private static String firstText(String... values) {
        for (String value : values) {
            if (value != null && !value.isBlank()) {
                return value;
            }
        }
        return null;
    }

    private static String requireText(String value, String name) {
        if (value == null || value.isBlank()) {
            throw new IllegalArgumentException(name + " cannot be blank");
        }
        return value;
    }
}
