package tech.easyflow.admin.service.ai;

import com.alibaba.fastjson.JSON;
import com.easyagents.flow.core.chain.Chain;
import com.easyagents.flow.core.chain.ChainConsts;
import com.easyagents.flow.core.chain.ChainStatus;
import com.easyagents.flow.core.chain.Edge;
import com.easyagents.flow.core.chain.Event;
import com.easyagents.flow.core.chain.Node;
import com.easyagents.flow.core.chain.event.ChainStatusChangeEvent;
import com.easyagents.flow.core.chain.event.EdgeConditionCheckFailedEvent;
import com.easyagents.flow.core.chain.event.EdgeTriggerEvent;
import com.easyagents.flow.core.chain.event.NodeEndEvent;
import com.easyagents.flow.core.chain.event.NodeStartEvent;
import com.easyagents.flow.core.chain.runtime.ChainExecutor;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.stereotype.Service;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;

import javax.annotation.PostConstruct;
import java.io.IOException;
import java.util.LinkedHashMap;
import java.util.Map;
import java.util.Objects;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.concurrent.atomic.AtomicLong;

/**
 * 将工作流运行事件转换为对话页可消费的 SSE 事件流。
 */
@Service
public class WorkflowChatEventStream {

    private static final Logger log =
            LoggerFactory.getLogger(WorkflowChatEventStream.class);
    private static final long SSE_TIMEOUT_MILLIS = 30L * 60L * 1000L;

    private final ChainExecutor chainExecutor;
    private final Map<String, StreamSession> sessions =
            new ConcurrentHashMap<>();

    /**
     * 创建工作流对话事件流服务。
     *
     * @param chainExecutor 工作流执行器
     */
    public WorkflowChatEventStream(ChainExecutor chainExecutor) {
        this.chainExecutor = chainExecutor;
    }

    /**
     * 注册工作流全局事件监听器。
     */
    @PostConstruct
    public void registerListeners() {
        chainExecutor.addEventListener(this::onEvent);
        chainExecutor.addErrorListener(this::onChainError);
    }

    /**
     * 启动工作流并返回其 SSE 连接。
     *
     * @param definitionId 工作流定义 ID
     * @param variables 运行变量
     * @return SSE 连接
     */
    public SseEmitter start(String definitionId, Map<String, Object> variables) {
        SseEmitter emitter = new SseEmitter(SSE_TIMEOUT_MILLIS);
        StreamSession session = new StreamSession(emitter);
        emitter.onTimeout(() -> disconnect(session, "运行连接超时"));
        emitter.onError(error -> disconnect(session, "运行连接已断开"));
        emitter.onCompletion(() -> removeSession(session));

        try {
            chainExecutor.executeAsync(
                    definitionId,
                    variables,
                    executeId -> {
                        session.attach(executeId);
                        sessions.put(executeId, session);
                        session.send("execution_started", Map.of(
                                "executeId", executeId
                        ));
                    }
            );
        } catch (RuntimeException | Error error) {
            session.fail(error);
            throw error;
        }
        return emitter;
    }

    /**
     * 将工作流事件转发到对应执行流。
     *
     * @param event 工作流事件
     * @param chain 当前工作流
     */
    private void onEvent(Event event, Chain chain) {
        StreamSession session = findSession(chain);
        if (session == null) {
            return;
        }
        if (event instanceof NodeStartEvent nodeStartEvent) {
            session.onNodeStarted(chain, nodeStartEvent);
            return;
        }
        if (event instanceof NodeEndEvent nodeEndEvent) {
            session.onNodeFinished(chain, nodeEndEvent);
            return;
        }
        if (event instanceof EdgeTriggerEvent edgeTriggerEvent) {
            session.onEdgeTriggered(chain, edgeTriggerEvent);
            return;
        }
        if (event instanceof EdgeConditionCheckFailedEvent failedEvent) {
            session.onEdgeConditionFailed(chain, failedEvent);
            return;
        }
        if (event instanceof ChainStatusChangeEvent statusEvent
                && Objects.equals(chain.getStateInstanceId(), session.executeId)) {
            session.onStatusChanged(chain, statusEvent.getStatus());
        }
    }

    /**
     * 将链级异常发送到客户端。
     *
     * @param error 链级异常
     * @param chain 当前工作流
     */
    private void onChainError(Throwable error, Chain chain) {
        StreamSession session = findSession(chain);
        if (session != null
                && Objects.equals(chain.getStateInstanceId(), session.executeId)) {
            session.send("execution_error", Map.of(
                    "message", safeErrorMessage(error)
            ));
        }
    }

    /**
     * 查找顶级执行对应的事件流会话。
     *
     * @param chain 当前工作流
     * @return 流会话；不存在时为 {@code null}
     */
    private StreamSession findSession(Chain chain) {
        if (chain == null) {
            return null;
        }
        String auditInstanceId = chain.getAuditInstanceId();
        if (auditInstanceId != null && !auditInstanceId.isBlank()) {
            StreamSession session = sessions.get(auditInstanceId);
            if (session != null) {
                return session;
            }
        }
        return sessions.get(chain.getStateInstanceId());
    }

    /**
     * 处理 SSE 连接异常，并取消尚未结束的工作流。
     *
     * @param session 流会话
     * @param message 取消原因
     */
    private void disconnect(StreamSession session, String message) {
        if (session == null || session.terminal.get()) {
            return;
        }
        String executeId = session.executeId;
        removeSession(session);
        if (executeId != null) {
            chainExecutor.cancel(executeId, message);
        }
    }

    /**
     * 移除流会话。
     *
     * @param session 流会话
     */
    private void removeSession(StreamSession session) {
        if (session != null && session.executeId != null) {
            sessions.remove(session.executeId, session);
        }
    }

    /**
     * 读取适合返回给用户的异常信息。
     *
     * @param error 异常
     * @return 非空异常信息
     */
    private String safeErrorMessage(Throwable error) {
        if (error == null || error.getMessage() == null
                || error.getMessage().isBlank()) {
            return "工作流执行失败";
        }
        return error.getMessage();
    }

    /**
     * 去掉顶级工作流结果中的内部状态控制字段。
     *
     * @param result 顶级工作流执行结果
     * @return 可直接交给会话结果区展示的最终输出
     */
    static Map<String, Object> visibleFinalOutput(
            Map<String, Object> result
    ) {
        Map<String, Object> visible = new LinkedHashMap<>();
        if (result != null) {
            visible.putAll(result);
        }
        visible.remove(ChainConsts.CHAIN_STATE_STATUS_KEY);
        visible.remove(ChainConsts.CHAIN_STATE_MESSAGE_KEY);
        visible.remove(ChainConsts.NODE_STATE_STATUS_KEY);
        visible.remove(ChainConsts.SCHEDULE_NEXT_NODE_DISABLED_KEY);
        return visible;
    }

    /**
     * 单次工作流执行的 SSE 会话。
     */
    private final class StreamSession {

        private final SseEmitter emitter;
        private final AtomicLong sequence = new AtomicLong();
        private final AtomicBoolean terminal = new AtomicBoolean(false);
        private volatile String executeId;

        /**
         * 创建流会话。
         *
         * @param emitter SSE 发送器
         */
        private StreamSession(SseEmitter emitter) {
            this.emitter = emitter;
        }

        /**
         * 绑定执行实例。
         *
         * @param executeId 执行实例 ID
         */
        private void attach(String executeId) {
            this.executeId = executeId;
        }

        /**
         * 处理节点开始事件。
         *
         * @param chain 当前工作流
         * @param event 节点开始事件
         */
        private void onNodeStarted(Chain chain, NodeStartEvent event) {
            Node node = event.getNode();
            Map<String, Object> data = new LinkedHashMap<>();
            data.put("attemptKey", event.getExecutionAttemptKey());
            data.put("nodeClass", node.getClass().getSimpleName());
            data.put("chainInstanceId", chain.getStateInstanceId());
            data.put("startedAt", System.currentTimeMillis());
            data.put("input", resolveNodeInput(chain, node));
            send("node_started", nodePayload(node, data));
        }

        /**
         * 处理节点完成事件并更新运行详情。
         *
         * @param chain 当前工作流
         * @param event 节点完成事件
         */
        private void onNodeFinished(Chain chain, NodeEndEvent event) {
            Node node = event.getNode();
            Map<String, Object> data = new LinkedHashMap<>();
            data.put("attemptKey", event.getExecutionAttemptKey());
            data.put("status", event.getStatus() == null
                    ? null
                    : event.getStatus().name());
            data.put("chainInstanceId", chain.getStateInstanceId());
            data.put("finishedAt", System.currentTimeMillis());
            data.put("output", event.getResult() == null
                    ? Map.of()
                    : event.getResult());
            if (event.getError() != null) {
                data.put("error", safeErrorMessage(event.getError()));
            }
            send("node_finished", nodePayload(node, data));
        }

        /**
         * 记录命中条件的流转分支。
         *
         * @param chain 当前工作流
         * @param event 边触发事件
         */
        private void onEdgeTriggered(Chain chain, EdgeTriggerEvent event) {
            if (event.getTrigger() == null) {
                return;
            }
            Edge edge = chain.getDefinition().getEdgeById(
                    event.getTrigger().getEdgeId());
            if (edge == null || edge.getCondition() == null) {
                return;
            }
            Node sourceNode = chain.getDefinition().getNodeById(
                    edge.getSource());
            sendEdgeTrace(chain, sourceNode, edge, "matched");
        }

        /**
         * 记录未命中条件的流转分支。
         *
         * @param chain 当前工作流
         * @param event 条件未命中事件
         */
        private void onEdgeConditionFailed(
                Chain chain,
                EdgeConditionCheckFailedEvent event
        ) {
            sendEdgeTrace(chain, event.getNode(), event.getEdge(), "skipped");
        }

        /**
         * 发送节点的条件判断轨迹。
         *
         * @param chain 当前工作流
         * @param sourceNode 条件来源节点
         * @param edge 被判断的边
         * @param outcome 判断结果
         */
        private void sendEdgeTrace(
                Chain chain,
                Node sourceNode,
                Edge edge,
                String outcome
        ) {
            if (sourceNode == null || edge == null) {
                return;
            }
            Node targetNode = chain.getDefinition().getNodeById(
                    edge.getTarget());
            Map<String, Object> data = new LinkedHashMap<>();
            data.put(
                    "attemptKey",
                    chain.currentExecutionAttemptKey(sourceNode.getId()));
            data.put("kind", "condition");
            data.put("outcome", outcome);
            data.put("edgeId", edge.getId());
            data.put("targetNodeId", edge.getTarget());
            data.put(
                    "targetNodeName",
                    targetNode == null ? edge.getTarget() : targetNode.getName());
            send("node_trace", nodePayload(sourceNode, data));
        }

        /**
         * 解析节点本次执行实际使用的输入。
         *
         * @param chain 当前工作流
         * @param node 当前节点
         * @return 可序列化的节点输入
         */
        private Map<String, Object> resolveNodeInput(Chain chain, Node node) {
            try {
                return chain.getExecutionState()
                        .resolveParametersPreservingReferences(node);
            } catch (RuntimeException error) {
                log.warn(
                        "Failed to resolve workflow node input, "
                                + "executeId={}, nodeId={}",
                        executeId,
                        node == null ? null : node.getId(),
                        error
                );
                return Map.of();
            }
        }

        /**
         * 处理工作流状态变化。
         *
         * @param chain 当前工作流
         * @param status 新状态
         */
        private void onStatusChanged(Chain chain, ChainStatus status) {
            if (status == null) {
                return;
            }
            if (status == ChainStatus.SUSPEND) {
                Map<String, Object> data = new LinkedHashMap<>();
                data.put("message", chain.getState().getMessage());
                data.put(
                        "parameters",
                        chain.getState().getSuspendForParameters()
                );
                send("execution_waiting", data);
                return;
            }
            if (!status.isTerminal()) {
                send("execution_status", Map.of("status", status.name()));
                return;
            }
            if (!terminal.compareAndSet(false, true)) {
                return;
            }
            String eventType = switch (status) {
                case SUCCEEDED -> "execution_finished";
                case CANCELLED -> "execution_cancelled";
                default -> "execution_failed";
            };
            Map<String, Object> data = new LinkedHashMap<>();
            data.put("status", status.name());
            data.put("message", chain.getState().getMessage());
            if (status == ChainStatus.SUCCEEDED) {
                data.put(
                        "output",
                        visibleFinalOutput(chain.getState().getExecuteResult())
                );
            }
            send(eventType, data);
            removeSession(this);
            emitter.complete();
        }

        /**
         * 发送 SSE 事件。
         *
         * @param type 事件类型
         * @param data 事件数据
         */
        private void send(String type, Map<String, ?> data) {
            long nextSequence = sequence.incrementAndGet();
            Map<String, Object> payload = new LinkedHashMap<>();
            payload.put("eventId", executeId + ":" + nextSequence);
            payload.put("sequence", nextSequence);
            payload.put("executeId", executeId);
            payload.put("type", type);
            payload.put("data", data);
            try {
                emitter.send(SseEmitter.event()
                        .id(String.valueOf(nextSequence))
                        .name("workflow")
                        // 显式发送 JSON 文本，避免全局 CBOR 转换器将 SSE data 编码为二进制。
                        .data(JSON.toJSONString(payload)));
            } catch (IOException | IllegalStateException error) {
                log.debug(
                        "workflow chat stream disconnected, executeId={}",
                        executeId,
                        error
                );
                disconnect(this, "运行连接已断开");
            }
        }

        /**
         * 在启动失败时关闭 SSE 会话。
         *
         * @param error 启动异常
         */
        private void fail(Throwable error) {
            if (terminal.compareAndSet(false, true)) {
                send("execution_failed", Map.of(
                        "message", safeErrorMessage(error)
                ));
                removeSession(this);
                emitter.completeWithError(error);
            }
        }

        /**
         * 构建带节点信息的事件数据。
         *
         * @param node 当前节点
         * @param values 业务数据
         * @return 事件数据
         */
        private Map<String, Object> nodePayload(
                Node node,
                Map<String, ?> values
        ) {
            Map<String, Object> payload = new LinkedHashMap<>();
            payload.put("nodeId", node == null ? null : node.getId());
            payload.put("nodeName", node == null ? null : node.getName());
            if (values != null) {
                payload.putAll(values);
            }
            return payload;
        }

    }
}
