/*
 * Copyright 2017-2025 noear.org and authors
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *      https://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package org.noear.solon.ai.chat.dialect;

import org.noear.snack4.Feature;
import org.noear.snack4.ONode;
import org.noear.snack4.Options;
import org.noear.snack4.json.JsonReader;
import org.noear.snack4.json.util.FormatUtil;
import org.noear.solon.Utils;
import org.noear.solon.ai.chat.content.ContentBlock;
import org.noear.solon.ai.chat.content.AudioBlock;
import org.noear.solon.ai.chat.*;
import org.noear.solon.ai.chat.tool.*;
import org.noear.solon.ai.chat.message.*;
import org.noear.solon.ai.chat.content.ImageBlock;
import org.noear.solon.ai.chat.content.TextBlock;
import org.noear.solon.ai.chat.content.VideoBlock;
import org.noear.solon.core.util.Assert;
import org.noear.solon.net.http.HttpUtils;
import org.noear.solon.net.http.impl.HttpSslSupplierAny;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.util.*;

/**
 * 聊天模型方言虚拟类
 *
 * @author noear
 * @since 3.1
 */
public abstract class AbstractChatDialect implements ChatDialect {
    private static final Logger LOG = LoggerFactory.getLogger(AbstractChatDialect.class);

    protected String getApiUrl(ChatConfig config){
        return config.getApiUrl();
    }

    public HttpUtils createHttpUtils(ChatConfig config, boolean isStream) {
        HttpUtils httpUtils = HttpUtils.http(getApiUrl(config))
                .ssl(HttpSslSupplierAny.getInstance())
                .timeout((int) config.getTimeout().getSeconds());

        if (config.getProxy() != null) {
            httpUtils.proxy(config.getProxy());
        }

        if (Utils.isNotEmpty(config.getApiKey())) {
            httpUtils.header("Authorization", "Bearer " + config.getApiKey());
        }

        if (Utils.isNotEmpty(config.getUserAgent())) {
            httpUtils.userAgent(config.getUserAgent());
        }

        httpUtils.headers(config.getHeaders());

        return httpUtils;
    }

    @Override
    public void prepareOutputSchemaInstruction(String outputSchema, StringBuilder instructionBuilder) {
        instructionBuilder.append("\n\n## [IMPORTANT: OUTPUT FORMAT]\n")
                .append("Format your response as a JSON object strictly following this schema:\n")
                .append("<output_schema>\n").append(outputSchema).append("\n</output_schema>\n")
                .append("Output only the raw JSON, beginning with '{' and ending with '}'.");
    }

    @Override
    public void prepareOutputFormatOptions(ChatOptions options) {
        options.optionSet("response_format", Utils.asMap("type", "json_object"));
    }

    protected void buildAssistantMessageNodeDo(ChatConfig config, ONode oNode, AssistantMessage msg) {
        oNode.set("role", msg.getRole().name().toLowerCase());

        if (Utils.isNotEmpty(msg.getResultContent())) {
            oNode.set("content", msg.getResultContent());
        }

        //兼容 r1 的 tool-call(可以再优化，只在最后一条加)
        if (Utils.isNotEmpty(msg.getReasoningFieldName())) {
            oNode.set(msg.getReasoningFieldName(), msg.getReasoning());
        } else if (Utils.isNotEmpty(config.getReasoningFieldName())) {
            oNode.set(config.getReasoningFieldName(), "");
        }

        if (Utils.isNotEmpty(msg.getToolCallsRaw())) {
            oNode.set("tool_calls", ONode.ofBean(msg.getToolCallsRaw()));
        }
    }

    protected void buildSystemMessageNodeDo(ChatConfig config, ONode oNode, SystemMessage msg) {
        oNode.set("role", msg.getRole().name().toLowerCase());
        oNode.set("content", msg.getContent());
    }

    protected void buildToolMessageNodeDo(ChatConfig config, ONode oNode, ToolMessage msg) {
        oNode.set("role", msg.getRole().name().toLowerCase());

        if (Utils.isNotEmpty(msg.getName())) {
            oNode.set("name", msg.getName());
        }

        if (Utils.isNotEmpty(msg.getToolCallId())) {
            oNode.set("tool_call_id", msg.getToolCallId());
        }

        if (msg.isMultiModal() == false) {
            oNode.set("content", msg.getContent());
        } else {
            oNode.getOrNew("content").then(n1 -> {
                for (ContentBlock m1 : msg.getBlocks()) {
                    ONode m1Node = null;

                    if (m1 instanceof TextBlock) {
                        TextBlock m1Text = (TextBlock) m1;
                        n1.addNew().set("type", "text").set("text", m1Text.getContent());
                    } else if (m1 instanceof ImageBlock) {
                        m1Node = n1.addNew();

                        m1Node.set("type", "image_url");
                        m1Node.getOrNew("image_url").set("url", m1.toDataString(true));

                    } else if (m1 instanceof AudioBlock) {
                        m1Node = n1.addNew();

                        m1Node.set("type", "audio_url");
                        m1Node.getOrNew("audio_url").set("url", m1.toDataString(true));
                    } else if (m1 instanceof VideoBlock) {
                        m1Node = n1.addNew();

                        m1Node.set("type", "video_url");
                        m1Node.getOrNew("video_url").set("url", m1.toDataString(true));
                    }

                    if (m1Node != null) {
                        if (Utils.isNotEmpty(m1.metas())) {
                            for (Map.Entry<String, Object> entry : m1.metas().entrySet()) {
                                if (m1Node.hasKey(entry.getKey()) == false) {
                                    m1Node.set(entry.getKey(), ONode.ofBean(entry.getValue()));
                                }
                            }
                        }
                    }
                }
            });
        }
    }

    protected void buildUserMessageNodeDo(ChatConfig config, ONode oNode, UserMessage msg) {
        oNode.set("role", msg.getRole().name().toLowerCase());

        if (msg.isMultiModal() == false) {
            //单模态
            oNode.set("content", msg.getContent());
        } else {
            //多模态
            oNode.getOrNew("content").then(n1 -> {
                for (ContentBlock block1 : msg.getBlocks()) {
                    ONode oNode1 = null;

                    if (block1 instanceof TextBlock) {
                        TextBlock m1Text = (TextBlock) block1;
                        n1.addNew().set("type", "text").set("text", m1Text.getContent());
                    } else if (block1 instanceof ImageBlock) {
                        oNode1 = n1.addNew();

                        oNode1.set("type", "image_url");
                        oNode1.getOrNew("image_url").set("url", block1.toDataString(true));

                    } else if (block1 instanceof AudioBlock) {
                        oNode1 = n1.addNew();

                        oNode1.set("type", "audio_url");
                        oNode1.getOrNew("audio_url").set("url", block1.toDataString(true));
                    } else if (block1 instanceof VideoBlock) {
                        oNode1 = n1.addNew();

                        oNode1.set("type", "video_url");
                        oNode1.getOrNew("video_url").set("url", block1.toDataString(true));
                    }

                    if (oNode1 != null) {
                        if (Utils.isNotEmpty(block1.metas())) {
                            for (Map.Entry<String, Object> entry : block1.metas().entrySet()) {
                                if (oNode1.hasKey(entry.getKey()) == false) {
                                    oNode1.set(entry.getKey(), ONode.ofBean(entry.getValue()));
                                }
                            }
                        }
                    }
                }
            });
        }
    }


    public ONode buildChatMessageNode(ChatConfig config, ChatMessage chatMessage) {
        ONode oNode = new ONode();
        if (chatMessage instanceof AssistantMessage) {
            buildAssistantMessageNodeDo(config, oNode, (AssistantMessage) chatMessage);
        } else if (chatMessage instanceof SystemMessage) {
            buildSystemMessageNodeDo(config, oNode, (SystemMessage) chatMessage);
        } else if (chatMessage instanceof ToolMessage) {
            buildToolMessageNodeDo(config, oNode, (ToolMessage) chatMessage);
        } else if (chatMessage instanceof UserMessage) {
            buildUserMessageNodeDo(config, oNode, (UserMessage) chatMessage);
        } else {
            throw new IllegalArgumentException("Unsupported chat message type: " + chatMessage.getClass());
        }

        return oNode;
    }

    /**
     * 构建请求工具节点
     */
    protected void buildReqToolsNode(ONode n, ChatConfig config, ChatOptions options, ChatMessage lastMessage) {
        //buildReqToolsNodeDo(n, config.getDefaultTools());
        buildReqToolsNodeDo(n, options.tools());
    }

    protected void buildReqToolsNodeDo(ONode n, Collection<FunctionTool> tools) {
        if (Utils.isEmpty(tools)) {
            return;
        }

        n.getOrNew("tools").then(n1 -> {
            for (FunctionTool func : tools) {
                n1.addNew().then(n2 -> {
                    n2.set("type", "function");
                    n2.getOrNew("function").then(toolNode -> {
                        toolNode.set("name", func.name());
                        toolNode.set("description", func.descriptionAndMeta());
                        toolNode.set("parameters", ONode.ofJson(func.inputSchema()));
                    });
                });
            }
        });
    }

    @Override
    public ONode buildRequestJson(ChatConfig config, ChatOptions options, List<ChatMessage> messages, boolean isStream) {
        return new ONode().then(n -> {
            if (Utils.isNotEmpty(config.getModel())) {
                n.set("model", config.getModel());
            }

            n.getOrNew("messages").then(n1 -> {
                for (ChatMessage m1 : messages) {
                    if (m1.isThinking() == false || m1.isToolCalls()) {
                        n1.add(buildChatMessageNode(config, m1));
                    }
                }
            });

            n.set("stream", isStream);

            for (Map.Entry<String, Object> kv : options.options().entrySet()) {
                n.set(kv.getKey(), ONode.ofBean(kv.getValue()));
            }

            ChatMessage lastMessage = messages.get(messages.size() - 1);
            buildReqToolsNode(n, config, options, lastMessage);
        });
    }

    @Override
    public ONode buildAssistantToolCallMessageNode(ChatResponseDefault resp, Map<String, ToolCallBuilder> toolCallBuilders) {
        ONode oNode = new ONode();
        oNode.set("role", "assistant");
        oNode.set("content", resp.getAggregationContent());
        oNode.getOrNew("tool_calls").asArray().then(n1 -> {
            for (Map.Entry<String, ToolCallBuilder> kv : toolCallBuilders.entrySet()) {
                //有可能没有
                n1.addNew().set("id", kv.getValue().idBuilder.toString())
                        .set("type", "function")
                        .getOrNew("function").then(n2 -> {
                            n2.set("name", kv.getValue().nameBuilder.toString());
                            if (kv.getValue().argumentsBuilder.length() > 0) {
                                n2.set("arguments", kv.getValue().argumentsBuilder.toString());
                            } else {
                                // vllm 不能传空
                                n2.set("arguments", "{}");
                            }
                        });
            }
        });

        return oNode;
    }

    @Override
    public AssistantMessage buildAssistantMessageByToolMessages(AssistantMessage toolCallMessage, List<ToolMessage> toolMessages) {
        //要求直接返回（转为新的响应消息）
        StringBuffer buf = new StringBuffer();
        for (ToolMessage toolMessage : toolMessages) {
            if (buf.length() > 0) {
                buf.append('\n');
            }
            buf.append(toolMessage.getContent());
        }

        AssistantMessage assistantMessage = ChatMessage.ofAssistant(buf.toString())
                .addMetadata("reason", "tool")
                .addMetadata("source", toolCallMessage.getResultContent());

        for (ToolMessage toolMessage : toolMessages) {
            assistantMessage.addMetadata(toolMessage.getMetadata());
        }

        return assistantMessage;
    }

    /**
     * 解析工具调用
     */
    protected List<ToolCall> parseToolCalls(ChatResponseDefault resp, ONode toolCallsNode) {
        if (toolCallsNode == null) {
            return null;
        }

        List<ToolCall> toolCalls = new ArrayList<>();

        for (ONode n1 : toolCallsNode.getArray()) {
            toolCalls.add(parseToolCall(resp, n1));
        }

        return toolCalls;
    }

    protected ToolCall parseToolCall(ChatResponseDefault resp, ONode n1) {
        String callId = n1.get("id").getString();

        if (Utils.isNotEmpty(callId)) {
            resp.lastToolCallId = callId;
        }

        String index = resp.lastToolCallId;

        ONode n1f = n1.get("function");
        String name = n1f.get("name").getString(); //可能是空的
        ONode n1fArgs = n1f.get("arguments");
        String argStr = n1fArgs.getString();

        if (n1fArgs.isString()) {
            //有可能是 json string（还可能只是流的中间消息）
            if (hasNestedJsonBlock(argStr)) {
                JsonReader reader = new JsonReader(argStr, Options.of(Feature.Read_AutoRepair));
                n1fArgs = reader.readLast();

                if (n1fArgs == null) {
                    LOG.warn("Parse tool arguments failed: {}", argStr);
                }
            }
        }

        Map<String, Object> argMap = new HashMap<>();
        if (n1fArgs != null) {
            if (n1fArgs.isObject()) {
                argMap = n1fArgs.toBean(Map.class);
            }
        }

        return new ToolCall(index, callId, name, argStr, argMap);
    }

    protected String parseAssistantMessageContent(ChatResponseDefault resp, ONode oContent) {
        if (oContent.isValue()) {
            //一般输出都是单值
            return oContent.getValueAs();
        } else {
            ONode contentItem = null;
            if (oContent.isArray()) {
                //有些输出会是列表（取第一个）
                if (oContent.getArrayUnsafe().size() > 0) {
                    contentItem = oContent.get(0);
                }
            } else if (oContent.isObject()) {
                //有些输出会是字典
                contentItem = oContent;
            }

            if (contentItem != null) {
                if (contentItem.isObject()) {
                    //优先取文本
                    if (contentItem.hasKey("text")) {
                        return contentItem.get("text").getValueAs();
                    } else if (contentItem.hasKey("image")) {
                        return contentItem.get("image").getValueAs();
                    } else if (contentItem.hasKey("audio")) {
                        return contentItem.get("audio").getValueAs();
                    } else if (contentItem.hasKey("video")) {
                        return contentItem.get("video").getValueAs();
                    }
                } else if (contentItem.isValue()) {
                    return contentItem.getValueAs();
                }
            }
        }

        return null;
    }

    public List<AssistantMessage> parseAssistantMessage(ChatResponseDefault resp, ONode oMessage) {
        List<AssistantMessage> messageList = new ArrayList<>();

        ONode oContent = oMessage.get("content");

        String content = parseAssistantMessageContent(resp, oContent);
        ONode toolCallsNode = oMessage.getOrNull("tool_calls");
        ONode searchResultsNode = oMessage.getOrNull("search_results");

        List<Map> toolCallsRaw = null;
        List<ToolCall> toolCalls = parseToolCalls(resp, toolCallsNode);
        List<Map> searchResultsRaw = null;

        if (Utils.isNotEmpty(toolCalls)) {
            toolCallsRaw = toolCallsNode.toBean(List.class);
            if (resp.in_thinking && resp.isStream()) {
                //说明是思考结束立刻调用了工具，需要添加思考的结束标识
                messageList.add(new AssistantMessage("</think>", true).reasoningFieldName(resp.reasoning_field_name));
                messageList.add(new AssistantMessage("\n\n", false).reasoningFieldName(resp.reasoning_field_name));
            }
            resp.in_thinking = false; //重置状态
        }

        if (searchResultsNode != null) {
            searchResultsRaw = searchResultsNode.toBean(List.class);
        }

        /**
         * 情况：
         * 有可能一直有：reasoning_content 或 reasoning
         * 有可能时有时无：reasoning_content 或 reasoning
         * 有可能一直无：...
         * 也可能和内容都为空: ...
         * */

        final String reasoning_content;
        if (oMessage.hasKey("reasoning_content")) {
            reasoning_content = oMessage.get("reasoning_content").getValueAs();
            resp.reasoning_field_name = "reasoning_content";
        } else if (oMessage.hasKey("reasoning")) {
            reasoning_content = oMessage.get("reasoning").getValueAs();
            resp.reasoning_field_name = "reasoning";
        } else {
            reasoning_content = null;
        }

        if (Utils.isNotEmpty(reasoning_content)) {
            resp.has_reasoning_field = true;
            //有思考专属内容的协议
            if (resp.isStream()) {
                //如果是流返回（可能要拆成多条流消息）
                if (Utils.isEmpty(content)) {
                    if (resp.in_thinking == false) {
                        //说明是第一次
                        messageList.add(new AssistantMessage("<think>", true).reasoningFieldName(resp.reasoning_field_name));
                        messageList.add(new AssistantMessage("\n\n", true).reasoningFieldName(resp.reasoning_field_name));
                        if (Utils.isNotEmpty(reasoning_content)) {
                            content = reasoning_content;
                        }
                    } else {
                        content = reasoning_content;
                    }

                    resp.in_thinking = true;
                } else {
                    if (resp.in_thinking) {
                        //说明是最后一次
                        messageList.add(new AssistantMessage("</think>", true).reasoningFieldName(resp.reasoning_field_name));
                        messageList.add(new AssistantMessage("\n\n", false).reasoningFieldName(resp.reasoning_field_name));
                    }

                    resp.in_thinking = false;
                }
            } else {
                //如查是单次返回
                if (Utils.isNotEmpty(reasoning_content)) {
                    content = "<think>\n\n" + reasoning_content + "</think>\n\n" + content;
                }
            }
        } else if (Utils.isNotEmpty(content)) {
            if (resp.has_reasoning_field) { //有些情况，后面就没字段了
                //有推理字段的
                if (resp.in_thinking) {
                    if (resp.isStream()) {
                        //说明是最后一次
                        messageList.add(new AssistantMessage("</think>", true).reasoningFieldName(resp.reasoning_field_name));
                        messageList.add(new AssistantMessage("\n\n", false).reasoningFieldName(resp.reasoning_field_name));
                    }

                    resp.in_thinking = false;
                }
            } else {
                //分析 think 状态（无推理字段的）
                if (resp.isStream()) {
                    //如果是流返回
                    if (content.startsWith("<think>")) {
                        resp.in_thinking = true;
                    } else {
                        if (resp.in_thinking) {
                            int thinkEnd = content.indexOf("</think>");
                            if (thinkEnd >= 0) { //可能是个开始符
                                resp.in_thinking = false;
                                messageList.add(new AssistantMessage(content, true).reasoningFieldName(resp.reasoning_field_name));
                                return messageList;
                            }
                        }
                    }
                }
            }
        }

        if (content != null || toolCallsRaw != null) {
            Object contentRaw = oContent.toBean();
            AssistantMessage message = new AssistantMessage(content, resp.in_thinking, contentRaw, toolCallsRaw, toolCalls, searchResultsRaw)
                    .reasoningFieldName(resp.reasoning_field_name);

            messageList.add(message);
        }

        return messageList;
    }


    protected boolean hasNestedJsonBlock(String str) {
        return FormatUtil.hasNestedJsonBlock(str);
    }
}