package tech.easyflow.ai.easyagents.listener;

import com.easyagents.core.message.AiMessage;
import com.easyagents.core.message.ToolCall;
import com.easyagents.core.message.ToolMessage;
import com.easyagents.core.model.chat.ChatModel;
import com.easyagents.core.model.chat.ChatOptions;
import com.easyagents.core.model.chat.StreamResponseListener;
import com.easyagents.core.model.chat.response.AiMessageResponse;
import com.easyagents.core.model.client.StreamContext;
import com.easyagents.core.prompt.MemoryPrompt;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import tech.easyflow.common.util.StringUtil;
import tech.easyflow.core.chat.protocol.ChatDomain;
import tech.easyflow.core.chat.protocol.ChatEnvelope;
import tech.easyflow.core.chat.protocol.ChatType;
import tech.easyflow.core.chat.protocol.MessageRole;
import tech.easyflow.core.chat.protocol.payload.ErrorPayload;
import tech.easyflow.core.chat.protocol.sse.ChatSseEmitter;
import tech.easyflow.core.runtime.ChatAssistantAccumulator;
import tech.easyflow.core.runtime.ChatRuntimeContext;
import tech.easyflow.core.runtime.ChatRuntimeExtKeys;
import tech.easyflow.core.runtime.ChatRuntimeManager;
import tech.easyflow.core.runtime.ChatRuntimeMessage;
import tech.easyflow.core.runtime.LegacyThinkingTagParser;

import java.math.BigInteger;
import java.util.Date;
import java.util.LinkedHashMap;
import java.util.List;
import java.util.Map;
import java.util.concurrent.atomic.AtomicBoolean;

public class ChatStreamListener implements StreamResponseListener {

    private static final Logger LOG = LoggerFactory.getLogger(ChatStreamListener.class);

    private final String conversationId;
    private final ChatModel chatModel;
    private final MemoryPrompt memoryPrompt;
    private final ChatSseEmitter sseEmitter;
    private final ChatOptions chatOptions;
    private final ChatRuntimeManager chatRuntimeManager;
    private final ChatRuntimeContext runtimeContext;
    private final ChatAssistantAccumulator assistantAccumulator;
    private final LegacyThinkingTagParser legacyThinkingTagParser = new LegacyThinkingTagParser();
    // 核心标记：是否允许执行onStop业务逻辑（仅最后一次无后续工具调用时为true）
    private boolean canStop = true;
    // 辅助标记：是否进入过工具调用（避免重复递归判断）
    private boolean hasToolCall = false;
    // 流式响应只能结束一次，避免重复发送导致 IllegalStateException
    private final AtomicBoolean completed = new AtomicBoolean(false);

    public ChatStreamListener(String conversationId, ChatModel chatModel, MemoryPrompt memoryPrompt, ChatSseEmitter sseEmitter,
                              ChatOptions chatOptions, ChatRuntimeManager chatRuntimeManager,
                              ChatRuntimeContext runtimeContext, ChatAssistantAccumulator assistantAccumulator) {
        this.conversationId = conversationId;
        this.chatModel = chatModel;
        this.memoryPrompt = memoryPrompt;
        this.sseEmitter = sseEmitter;
        this.chatOptions = chatOptions;
        this.chatRuntimeManager = chatRuntimeManager;
        this.runtimeContext = runtimeContext;
        this.assistantAccumulator = assistantAccumulator;
    }

    @Override
    public void onStart(StreamContext context) {
        legacyThinkingTagParser.reset();
        StreamResponseListener.super.onStart(context);
    }

    @Override
    public void onMessage(StreamContext context, AiMessageResponse aiMessageResponse) {
        try {
            if (completed.get() || sseEmitter.isClosed()) {
                stopStreamClient(context, "emitter_closed_before_message", null);
                return;
            }
            AiMessage aiMessage = aiMessageResponse.getMessage();
            if (aiMessage == null) {
                return;
            }
            if (aiMessage.isFinalDelta() && aiMessageResponse.hasToolCalls()) {
                flushLegacyThinkingSegments();
                this.canStop = false; // 工具调用期间，禁止执行onStop
                this.hasToolCall = true; // 标记已进入过工具调用
                List<ToolCall> toolCalls = aiMessage.getToolCalls();
                if (toolCalls != null) {
                    for (ToolCall toolCall : toolCalls) {
                        assistantAccumulator.appendToolCall(toolCall.getId(), toolCall.getName(), toolCall.getArguments());
                        sendToolCallEnvelope(toolCall);
                    }
                }
                applyToolCallHistorySnapshot(aiMessage);
                aiMessage.setContent(null);
                memoryPrompt.addMessage(aiMessage);
                List<ToolMessage> toolMessages = aiMessageResponse.executeToolCallsAndGetToolMessages();
                for (ToolMessage toolMessage : toolMessages) {
                    memoryPrompt.addMessage(toolMessage);
                    assistantAccumulator.appendToolResult(toolMessage.getToolCallId(), null, toolMessage.getContent());
                    sendToolResultEnvelope(toolMessage);
                }
                chatModel.chatStream(memoryPrompt, this, chatOptions);
            } else {
                if (this.hasToolCall) {
                    this.canStop = true;
                }
                handleAssistantDelta(aiMessage);
            }
        } catch (Exception e) {
            LOG.error("Chat stream onMessage failed, conversationId={}, message={}, exception={}",
                    conversationId, e.getMessage(), e.toString(), e);
            if (completed.compareAndSet(false, true)) {
                sendSystemError(sseEmitter, "流式响应异常，请重试", e);
            }
            stopStreamClient(context, "on_message_exception", e);
        }
    }

    @Override
    public void onStop(StreamContext context) {
        // 仅当canStop为true（最后一次无后续工具调用的响应）时，执行业务逻辑
        if (this.canStop && completed.compareAndSet(false, true)) {
            if (context.getThrowable() != null) {
                chatRuntimeManager.recordFailure(runtimeContext, context.getThrowable());
                sendSystemError(sseEmitter, context.getThrowable().getMessage(), context.getThrowable());
                return;
            }
            flushLegacyThinkingSegments();
            memoryPrompt.addMessage(context.getFullMessage());
            chatRuntimeManager.recordAssistantCompleted(runtimeContext, buildAssistantCompletedMessage(context));
            chatRuntimeManager.recordCompleted(runtimeContext);
            ChatEnvelope<Map<String, String>> chatEnvelope = new ChatEnvelope<>();
            chatEnvelope.setDomain(ChatDomain.SYSTEM);
            boolean doneSent = sseEmitter.sendDone(chatEnvelope);
            if (!doneSent) {
                LOG.warn("sendDone skipped because emitter is closed, conversationId={}", conversationId);
            }
            StreamResponseListener.super.onStop(context);
        }

    }

    @Override
    public void onFailure(StreamContext context, Throwable throwable) {
        if (throwable != null) {
            LOG.error("Chat stream onFailure, conversationId={}, message={}, exception={}",
                    conversationId, throwable.getMessage(), throwable.toString(), throwable);
        }
        if (throwable != null && completed.compareAndSet(false, true)) {
            chatRuntimeManager.recordFailure(runtimeContext, throwable);
            sendSystemError(sseEmitter, throwable.getMessage(), throwable);
        }
        stopStreamClient(context, "on_failure", throwable);
    }

    private void sendChatEnvelope(ChatSseEmitter sseEmitter, String deltaContent, ChatType chatType) {
        if (deltaContent == null || deltaContent.isEmpty()) {
            return;
        }

        ChatEnvelope<Map<String, String>> chatEnvelope = new ChatEnvelope<>();
        chatEnvelope.setDomain(ChatDomain.LLM);
        chatEnvelope.setType(chatType);

        Map<String, String> deltaMap = new LinkedHashMap<>();
        deltaMap.put("conversation_id", this.conversationId);
        deltaMap.put("role", MessageRole.ASSISTANT.getValue());
        deltaMap.put("delta", deltaContent);
        chatEnvelope.setPayload(deltaMap);
        chatEnvelope.setMeta(buildStreamMeta());

        boolean sent = sseEmitter.send(chatEnvelope);
        if (!sent) {
            throw new IllegalStateException("SSE emitter has already completed while sending chat envelope");
        }
    }

    /**
     * 归一化并发送单个模型增量。
     *
     * @param aiMessage 模型增量消息
     */
    private void handleAssistantDelta(AiMessage aiMessage) {
        String reasoningContent = aiMessage.getReasoningContent();
        if (StringUtil.hasText(reasoningContent)) {
            emitAssistantSegments(legacyThinkingTagParser.acceptReasoning(reasoningContent));
            return;
        }
        emitAssistantSegments(legacyThinkingTagParser.acceptContent(aiMessage.getContent()));
    }

    /**
     * 发送兼容解析后的思考与正文片段。
     *
     * @param segments 解析片段
     */
    private void emitAssistantSegments(List<LegacyThinkingTagParser.Segment> segments) {
        for (LegacyThinkingTagParser.Segment segment : segments) {
            String text = segment.getText();
            if (segment.getType() == LegacyThinkingTagParser.SegmentType.REASONING) {
                assistantAccumulator.appendReasoning(text);
                chatRuntimeManager.recordAssistantDelta(runtimeContext,
                        buildAssistantDeltaMessage(text, ChatType.THINKING));
                sendChatEnvelope(sseEmitter, text, ChatType.THINKING);
                continue;
            }
            assistantAccumulator.appendContent(text);
            chatRuntimeManager.recordAssistantDelta(runtimeContext,
                    buildAssistantDeltaMessage(text, ChatType.MESSAGE));
            sendChatEnvelope(sseEmitter, text, ChatType.MESSAGE);
        }
    }

    /**
     * 收口并发送旧标签解析器仍保留的少量前缀。
     */
    private void flushLegacyThinkingSegments() {
        emitAssistantSegments(legacyThinkingTagParser.finish());
    }

    private void sendToolCallEnvelope(ToolCall toolCall) {
        if (toolCall == null) {
            return;
        }
        ChatEnvelope<Map<String, Object>> chatEnvelope = new ChatEnvelope<>();
        chatEnvelope.setDomain(ChatDomain.TOOL);
        chatEnvelope.setType(ChatType.TOOL_CALL);

        Map<String, Object> payload = new LinkedHashMap<>();
        payload.put("tool_call_id", toolCall.getId());
        payload.put("name", toolCall.getName());
        payload.put("arguments", toolCall.getArguments());
        chatEnvelope.setPayload(payload);
        chatEnvelope.setMeta(buildStreamMeta());
        boolean sent = sseEmitter.send(chatEnvelope);
        if (!sent) {
            throw new IllegalStateException("SSE emitter has already completed while sending tool call envelope");
        }
    }

    private void sendToolResultEnvelope(ToolMessage toolMessage) {
        if (toolMessage == null) {
            return;
        }
        ChatEnvelope<Map<String, Object>> chatEnvelope = new ChatEnvelope<>();
        chatEnvelope.setDomain(ChatDomain.TOOL);
        chatEnvelope.setType(ChatType.TOOL_RESULT);

        Map<String, Object> payload = new LinkedHashMap<>();
        payload.put("tool_call_id", toolMessage.getToolCallId());
        payload.put("result", toolMessage.getContent());
        chatEnvelope.setPayload(payload);
        chatEnvelope.setMeta(buildStreamMeta());
        boolean sent = sseEmitter.send(chatEnvelope);
        if (!sent) {
            throw new IllegalStateException("SSE emitter has already completed while sending tool result envelope");
        }
    }

    public void sendSystemError(ChatSseEmitter sseEmitter,
                                String message,
                                Throwable throwable) {
        String errorMessage = StringUtil.hasText(message) ? message : "系统繁忙，请稍后重试";
        if (throwable != null) {
            LOG.error("sendSystemError, conversationId={}, message={}, exception={}",
                    conversationId, throwable.getMessage(), throwable.toString(), throwable);
        }
        ChatEnvelope<ErrorPayload> envelope = new ChatEnvelope<>();
        ErrorPayload payload = new ErrorPayload();
        payload.setMessage(errorMessage);
        payload.setCode("SYSTEM_ERROR");
        payload.setRetryable(false);
        envelope.setPayload(payload);
        envelope.setDomain(ChatDomain.SYSTEM);
        envelope.setType(ChatType.ERROR);
        envelope.setMeta(buildStreamMeta());
        boolean sent = sseEmitter.sendError(envelope);
        if (!sent) {
            LOG.warn("sendSystemError skipped because emitter is closed, conversationId={}", conversationId);
        }
        sseEmitter.complete();
    }

    private Map<String, Object> buildStreamMeta() {
        Map<String, Object> meta = new LinkedHashMap<>();
        BigInteger roundId = getBigIntegerExt(ChatRuntimeExtKeys.CURRENT_ROUND_ID);
        Integer roundNo = getIntegerExt(ChatRuntimeExtKeys.CURRENT_ROUND_NO);
        Integer variantIndex = getIntegerExt(ChatRuntimeExtKeys.CURRENT_VARIANT_INDEX);
        BigInteger regenerateRoundId = getBigIntegerExt(ChatRuntimeExtKeys.REGENERATE_ROUND_ID);
        if (roundId != null) {
            meta.put("roundId", roundId.toString());
        }
        if (roundNo != null) {
            meta.put("roundNo", roundNo);
        }
        if (variantIndex != null) {
            meta.put("variantIndex", variantIndex);
        }
        meta.put("regenerate", regenerateRoundId != null);
        if (regenerateRoundId != null) {
            meta.put("regenerateRoundId", regenerateRoundId.toString());
        }
        return meta;
    }

    private BigInteger getBigIntegerExt(String key) {
        Object value = runtimeContext == null || runtimeContext.getExt() == null
                ? null
                : runtimeContext.getExt().get(key);
        if (value == null) {
            return null;
        }
        if (value instanceof BigInteger number) {
            return number;
        }
        return new BigInteger(String.valueOf(value));
    }

    private Integer getIntegerExt(String key) {
        Object value = runtimeContext == null || runtimeContext.getExt() == null
                ? null
                : runtimeContext.getExt().get(key);
        if (value == null) {
            return null;
        }
        if (value instanceof Integer number) {
            return number;
        }
        return Integer.parseInt(String.valueOf(value));
    }

    private void stopStreamClient(StreamContext context, String reason, Throwable source) {
        try {
            if (context != null && context.getClient() != null) {
                context.getClient().stop();
                LOG.warn("Stopped stream client, conversationId={}, reason={}", conversationId, reason);
            }
        } catch (Exception stopEx) {
            LOG.error("Stop stream client failed, conversationId={}, reason={}, message={}, exception={}",
                    conversationId, reason, stopEx.getMessage(), stopEx.toString(), stopEx);
            if (source != null) {
                LOG.error("Original stream failure, conversationId={}, message={}, exception={}",
                        conversationId, source.getMessage(), source.toString(), source);
            }
        }
    }

    private ChatRuntimeMessage buildAssistantDeltaMessage(String delta, ChatType chatType) {
        ChatRuntimeMessage message = new ChatRuntimeMessage();
        message.setRole("assistant");
        message.setContentType(chatType == ChatType.THINKING ? "THINKING" : "TEXT");
        message.setContentText(delta);
        message.setCreatedAt(new Date());
        message.setSenderId(runtimeContext.getAssistantId());
        message.setSenderName(runtimeContext.getAssistantName());
        applyRoundMetadata(message);
        return message;
    }

    private ChatRuntimeMessage buildAssistantCompletedMessage(StreamContext context) {
        ChatRuntimeMessage message = new ChatRuntimeMessage();
        message.setRole("assistant");
        message.setContentType("TEXT");
        String fullContent = context != null && context.getFullMessage() != null ? context.getFullMessage().getContent() : null;
        String normalizedContent = assistantAccumulator.getContent();
        message.setContentText(legacyThinkingTagParser.isLegacyFormatDetected()
                ? normalizedContent
                : (StringUtil.hasText(fullContent) ? fullContent : normalizedContent));
        message.setContentPayload(assistantAccumulator.buildPayload(message.getContentText()));
        message.setCreatedAt(new Date());
        message.setSenderId(runtimeContext.getAssistantId());
        message.setSenderName(runtimeContext.getAssistantName());
        applyRoundMetadata(message);
        return message;
    }

    /**
     * 把当前轮次元数据写回 assistant 消息，确保 SSE 与后续持久化链路都能识别轮次和版本。
     *
     * @param message assistant 运行时消息
     */
    private void applyRoundMetadata(ChatRuntimeMessage message) {
        if (message == null) {
            return;
        }
        message.setRoundId(getBigIntegerExt(ChatRuntimeExtKeys.CURRENT_ROUND_ID));
        message.setRoundNo(getIntegerExt(ChatRuntimeExtKeys.CURRENT_ROUND_NO));
        message.setVariantIndex(getIntegerExt(ChatRuntimeExtKeys.CURRENT_VARIANT_INDEX));
    }

    /**
     * 在 tool call assistant 写入临时 memory 前，把 reasoning/content 快照回填到消息对象中，
     * 以便前端 history 透传和 DeepSeek 下一轮请求都能拿到完整链路。
     *
     * @param aiMessage tool call assistant 消息
     */
    private void applyToolCallHistorySnapshot(AiMessage aiMessage) {
        if (aiMessage == null) {
            return;
        }
        if (!StringUtil.hasText(aiMessage.getReasoningContent())) {
            aiMessage.setReasoningContent(assistantAccumulator.getLatestToolCallReasoning());
        }
        if (!StringUtil.hasText(aiMessage.getTextContent())) {
            aiMessage.setFullContent(assistantAccumulator.getLatestToolCallContent());
        }
    }

}
