package com.easyagents.flow.support.provider;

import com.easyagents.core.message.AiMessage;
import com.easyagents.core.message.SystemMessage;
import com.easyagents.core.model.chat.ChatModel;
import com.easyagents.core.model.chat.response.AiMessageResponse;
import com.easyagents.core.prompt.SimplePrompt;
import com.easyagents.flow.core.chain.Chain;
import com.easyagents.flow.core.llm.Llm;
import com.easyagents.flow.core.node.LlmNode;
import com.easyagents.flow.core.util.StringUtil;

import java.util.List;

public class EasyAgentsLlm implements Llm {

    private ChatModel chatModel;

    public ChatModel getChatModel() {
        return chatModel;
    }

    public void setChatModel(ChatModel chatModel) {
        this.chatModel = chatModel;
    }

    @Override
    public String chat(MessageInfo messageInfo, ChatOptions options, LlmNode llmNode, Chain chain) {

        SimplePrompt prompt = new SimplePrompt(messageInfo.getMessage());

        // 系统提示词
        if (StringUtil.hasText(messageInfo.getSystemMessage())) {
            prompt.setSystemMessage(SystemMessage.of(messageInfo.getSystemMessage()));
        }

        //  图片
        List<String> images = messageInfo.getImages();
        if (images != null && !images.isEmpty()) {
            for (String image : images) {
                prompt.addImageUrl(image);
            }
        }

        com.easyagents.core.model.chat.ChatOptions chatOptions = new com.easyagents.core.model.chat.ChatOptions();
        chatOptions.setSeed(options.getSeed());
        chatOptions.setTemperature(options.getTemperature());
        chatOptions.setTopP(options.getTopP());
        chatOptions.setTopK(options.getTopK());
        chatOptions.setMaxTokens(options.getMaxTokens());
        chatOptions.setStop(options.getStop());

        AiMessageResponse response = chatModel.chat(prompt, chatOptions);
        if (response == null) {
            throw new RuntimeException("EasyAgentsLlm can not get response!");
        }

        if (response.isError()) {
            throw new RuntimeException("EasyAgentsLlm error: " + response.getErrorMessage());
        }

        AiMessage aiMessage = response.getMessage();
        if (aiMessage != null) {
            return aiMessage.getTextContent();
        }

        throw new RuntimeException("EasyAgentsLlm can not get aiMessage!");
    }
}
