package tech.easyflow.core.runtime;

import java.util.ArrayList;
import java.util.LinkedHashMap;
import java.util.List;
import java.util.Map;

/**
 * 负责聚合单轮聊天中的 assistant thinking、tool 调用和最终回答，
 * 并产出可用于历史回放的结构化 payload。
 */
public class ChatAssistantAccumulator {

    private final StringBuilder content = new StringBuilder();
    private final StringBuilder reasoning = new StringBuilder();
    private final StringBuilder displayReasoning = new StringBuilder();
    private final List<Map<String, Object>> chains = new ArrayList<>();
    private final List<Map<String, Object>> messageChain = new ArrayList<>();
    private final List<Map<String, Object>> toolMessages = new ArrayList<>();
    private Map<String, Object> latestToolCallAssistant;
    private boolean toolCallBatchOpen;

    /**
     * 追加当前 assistant 片段的文本内容。
     *
     * @param delta 内容增量
     */
    public void appendContent(String delta) {
        if (delta != null && !delta.isEmpty()) {
            content.append(delta);
        }
    }

    /**
     * 追加当前 assistant 片段的 reasoning 内容。
     *
     * @param delta reasoning 增量
     */
    public void appendReasoning(String delta) {
        if (delta != null && !delta.isEmpty()) {
            reasoning.append(delta);
            displayReasoning.append(delta);
        }
    }

    /**
     * 记录 tool call，同时把当前 assistant 片段固化为一条结构化 assistant 消息。
     *
     * @param id tool call id
     * @param name tool 名称
     * @param arguments tool 参数
     */
    public void appendToolCall(String id, String name, Object arguments) {
        appendToolCall(id, name, null, arguments);
    }

    /**
     * 记录 tool call，同时保留面向前端展示的工具名称。
     *
     * @param id tool call id
     * @param name tool 名称
     * @param displayName tool 展示名称
     * @param arguments tool 参数
     */
    public void appendToolCall(String id, String name, String displayName, Object arguments) {
        Map<String, Object> chain = findToolChain(id, name);
        chain.put("status", "TOOL_CALL");
        chain.put("arguments", arguments);
        putIfNotBlank(chain, "toolDisplayName", displayName);

        Map<String, Object> assistantMessage = ensureToolCallAssistantMessage();
        @SuppressWarnings("unchecked")
        List<Map<String, Object>> toolCalls = (List<Map<String, Object>>) assistantMessage.computeIfAbsent("toolCalls",
                key -> new ArrayList<Map<String, Object>>());
        Map<String, Object> toolCall = new LinkedHashMap<>();
        toolCall.put("id", id);
        toolCall.put("name", name);
        toolCall.put("arguments", arguments == null ? null : String.valueOf(arguments));
        putIfNotBlank(toolCall, "toolDisplayName", displayName);
        toolCalls.add(toolCall);
    }

    /**
     * 记录 tool result，并附加到结构化消息链中。
     *
     * @param id tool call id
     * @param name tool 名称
     * @param result tool 结果
     */
    public void appendToolResult(String id, String name, Object result) {
        appendToolResult(id, name, null, result);
    }

    /**
     * 记录 tool result，并保留面向前端展示的工具名称。
     *
     * @param id tool call id
     * @param name tool 名称
     * @param displayName tool 展示名称
     * @param result tool 结果
     */
    public void appendToolResult(String id, String name, String displayName, Object result) {
        Map<String, Object> chain = findToolChain(id, name);
        chain.put("status", "TOOL_RESULT");
        chain.put("result", result);
        putIfNotBlank(chain, "toolDisplayName", displayName);
        Map<String, Object> toolMessage = ChatRuntimeHistoryPayloadHelper.toolMessage(
                id,
                result == null ? null : String.valueOf(result)
        );
        toolMessages.add(toolMessage);
        messageChain.add(ChatRuntimeHistoryPayloadHelper.deepCopyMap(toolMessage));
        toolCallBatchOpen = false;
    }

    /**
     * 获取当前 assistant 片段的文本内容。
     *
     * @return 文本内容
     */
    public String getContent() {
        return content.toString();
    }

    /**
     * 获取最近一次 tool-call assistant 的 reasoning 内容，供实时内存消息回写复用。
     *
     * @return reasoning 内容
     */
    public String getLatestToolCallReasoning() {
        return latestToolCallAssistant == null ? null : stringValue(latestToolCallAssistant.get("reasoningContent"));
    }

    /**
     * 获取最近一次 tool-call assistant 的内容，供实时内存消息回写复用。
     *
     * @return 内容
     */
    public String getLatestToolCallContent() {
        return latestToolCallAssistant == null ? null : stringValue(latestToolCallAssistant.get("content"));
    }

    /**
     * 产出结构化 payload。
     *
     * @param finalContent 最终 assistant 文本
     * @return payload
     */
    public Map<String, Object> buildPayload(String finalContent) {
        List<Map<String, Object>> payloadChains = new ArrayList<>();
        if (displayReasoning.length() > 0) {
            Map<String, Object> think = new LinkedHashMap<>();
            think.put("reasoning_content", displayReasoning.toString());
            think.put("thinkingStatus", "end");
            think.put("thinlCollapse", Boolean.TRUE);
            payloadChains.add(think);
        }
        payloadChains.addAll(chains);
        List<Map<String, Object>> payloadMessageChain = ChatRuntimeHistoryPayloadHelper.deepCopyList(messageChain);
        Map<String, Object> finalAssistantMessage = buildFinalAssistantMessage(finalContent);
        if (!finalAssistantMessage.isEmpty()) {
            payloadMessageChain.add(finalAssistantMessage);
        }
        return ChatRuntimeHistoryPayloadHelper.buildPayload(payloadMessageChain, toolMessages, payloadChains);
    }

    private Map<String, Object> findToolChain(String id, String name) {
        for (Map<String, Object> chain : chains) {
            if (String.valueOf(chain.get("id")).equals(id)) {
                if (name != null && !name.isEmpty()) {
                    chain.put("name", name);
                }
                return chain;
            }
        }
        Map<String, Object> chain = new LinkedHashMap<>();
        chain.put("id", id);
        chain.put("name", name);
        chains.add(chain);
        return chain;
    }

    private Map<String, Object> ensureToolCallAssistantMessage() {
        if (toolCallBatchOpen && latestToolCallAssistant != null && !hasPendingAssistantContent()) {
            return latestToolCallAssistant;
        }
        latestToolCallAssistant = ChatRuntimeHistoryPayloadHelper.assistantMessage(
                content.length() == 0 ? null : content.toString(),
                reasoning.length() == 0 ? null : reasoning.toString(),
                null
        );
        messageChain.add(latestToolCallAssistant);
        content.setLength(0);
        reasoning.setLength(0);
        toolCallBatchOpen = true;
        return latestToolCallAssistant;
    }

    private Map<String, Object> buildFinalAssistantMessage(String finalContent) {
        String assistantContent = finalContent;
        if ((assistantContent == null || assistantContent.isEmpty()) && content.length() > 0) {
            assistantContent = content.toString();
        }
        if ((assistantContent == null || assistantContent.isEmpty()) && reasoning.length() == 0) {
            return new LinkedHashMap<>();
        }
        return ChatRuntimeHistoryPayloadHelper.assistantMessage(
                assistantContent,
                reasoning.length() == 0 ? null : reasoning.toString(),
                null
        );
    }

    private boolean hasPendingAssistantContent() {
        return content.length() > 0 || reasoning.length() > 0;
    }

    private String stringValue(Object value) {
        return value == null ? null : String.valueOf(value);
    }

    private void putIfNotBlank(Map<String, Object> target, String key, String value) {
        if (value != null && !value.isBlank()) {
            target.put(key, value);
        }
    }
}
