/*
 * 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.llm.dialect.ollama;

import org.noear.snack4.Feature;
import org.noear.snack4.ONode;
import org.noear.snack4.Options;
import org.noear.snack4.json.JsonReader;
import org.noear.solon.Utils;
import org.noear.solon.ai.chat.content.ContentBlock;
import org.noear.solon.ai.AiUsage;
import org.noear.solon.ai.chat.content.AudioBlock;
import org.noear.solon.ai.chat.ChatChoice;
import org.noear.solon.ai.chat.ChatConfig;
import org.noear.solon.ai.chat.ChatException;
import org.noear.solon.ai.chat.ChatResponseDefault;
import org.noear.solon.ai.chat.dialect.AbstractChatDialect;
import org.noear.solon.ai.chat.message.AssistantMessage;
import org.noear.solon.ai.chat.message.UserMessage;
import org.noear.solon.ai.chat.tool.ToolCall;
import org.noear.solon.ai.chat.tool.ToolCallBuilder;
import org.noear.solon.ai.chat.content.ImageBlock;
import org.noear.solon.ai.chat.content.VideoBlock;
import org.noear.solon.core.util.Assert;
import org.noear.solon.core.util.DateUtil;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.util.Date;
import java.util.List;
import java.util.Map;

/**
 * Ollama 聊天模型方言
 *
 * @author noear
 * @since 3.1
 */
public class OllamaChatDialect extends AbstractChatDialect {

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

    private static OllamaChatDialect instance = new OllamaChatDialect();

    public static OllamaChatDialect getInstance() {
        return instance;
    }

    /**
     * 匹配检测
     *
     * @param config 聊天配置
     */
    @Override
    public boolean matched(ChatConfig config) {
        return "ollama".equalsIgnoreCase(config.getProvider()) ||
                (Assert.isEmpty(config.getProvider()) && config.getApiUrl().endsWith("/api/chat"));
    }

    @Override
    protected String getApiUrl(ChatConfig config) {

        //处理后缀#
        int index = config.getApiUrl().indexOf('#');
        if (index > 0) {
            return config.getApiUrl().substring(0, index);
        }

        //自动补全地址
        if (config.getApiUrl().endsWith("/api/chat")) {
            return config.getApiUrl();
        } else {
            if (config.getApiUrl().endsWith("/")) {
                return config.getApiUrl() + "api/chat";
            } else {
                return config.getApiUrl() + "/api/chat";
            }
        }
    }

    @Override
    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.set("content", msg.getContent());

            for (ContentBlock block1 : msg.getBlocks()) {
                if (block1 instanceof ImageBlock) {
                    oNode.getOrNew("images").add(block1.toDataString(false));
                } else if (block1 instanceof AudioBlock) {
                    oNode.getOrNew("audios").add(block1.toDataString(false));
                } else if (block1 instanceof VideoBlock) {
                    oNode.getOrNew("videos").add(block1.toDataString(false));
                }
            }
        }
    }

    @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());
                            n2.set("arguments", ONode.ofJson(kv.getValue().argumentsBuilder.toString()));
                        });
            }
        });

        return oNode;
    }

    @Override
    public boolean parseResponseJson(ChatConfig config, ChatResponseDefault resp, String json) {
        //解析
        ONode oResp = ONode.ofJson(json);

        if (oResp.isObject() == false) {
            return false;
        }

        if (oResp.hasKey("error")) {
            resp.setError(new ChatException(oResp.get("error").getString()));
        } else {
            resp.setModel(oResp.get("model").getString());
            resp.setFinished(oResp.get("done").getBoolean());
            String done_reason = oResp.get("done_reason").getString();

            String createdStr = oResp.get("created_at").getString();
            if (createdStr != null) {
                createdStr = createdStr.substring(0, createdStr.indexOf(".") + 4);
            }
            Date created = DateUtil.parseTry(createdStr);
            List<AssistantMessage> messageList = parseAssistantMessage(resp, oResp.get("message"));
            for (AssistantMessage msg1 : messageList) {
                resp.addChoice(new ChatChoice(0, created, done_reason, msg1));
            }

            if (Utils.isNotEmpty(done_reason)) {
                resp.lastFinishReason = done_reason;
            }

            if (resp.isFinished()) {
                long promptTokens = oResp.get("prompt_eval_count").getLong();
                long completionTokens = oResp.get("eval_count").getLong();
                long totalTokens = promptTokens + completionTokens;

                resp.setUsage(new AiUsage(promptTokens, 0L, completionTokens, totalTokens, oResp));

                if (resp.hasChoices() == false) {
                    resp.addChoice(new ChatChoice(0, created, resp.getLastFinishReasonNormalized(), new AssistantMessage("")));
                }
            }
        }

        return true;
    }

    @Override
    protected ToolCall parseToolCall(ChatResponseDefault resp, ONode n1) {
        String callId = n1.get("id").getString();//可能是空的

        ONode n1f = n1.get("function");
        String name = n1f.get("name").getString();
        ONode n1fArgs = n1f.get("arguments");
        String argStr = n1fArgs.getString();

        String index = name;

        if (n1fArgs.isValue()) {
            //有可能是 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 = null;
        if (n1fArgs != null) {
            if (n1fArgs.isObject()) {
                argMap = n1fArgs.toBean(Map.class);
            }
        }

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