package tech.easyflow.ai.easyagents.memory;

import com.easyagents.core.memory.DefaultChatMemory;
import com.easyagents.core.message.*;
import com.alibaba.fastjson.JSON;
import com.alibaba.fastjson.serializer.SerializerFeature;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import tech.easyflow.ai.entity.BotMessage;
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.sse.ChatSseEmitter;

import java.math.BigInteger;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.stream.Collectors;


public class DefaultBotMessageMemory extends DefaultChatMemory {

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

    private final ChatSseEmitter sseEmitter;

    private final List<Map<String, String>> messages;
    public DefaultBotMessageMemory(BigInteger conversationId, ChatSseEmitter sseEmitter, List<Map<String, String>> messages) {
        super(conversationId);
        this.sseEmitter = sseEmitter;
        this.messages = messages;
    }

    @Override
    public List<Message> getMessages(int count) {
        List<Message> list = new ArrayList<>(messages.size());
        for (Map<String, String> msg : messages) {
            BotMessage botMessage = new BotMessage();
            botMessage.setRole(msg.get("role"));
            botMessage.setContent(msg.get("content"));
            Message message = botMessage.getContentAsMessage();
            list.add(message);
        }
        List<Message> collect = list.stream()
                .limit(count)
                .collect(Collectors.toList());
        return collect;
    }

    @Override
    public void addMessage(Message message) {
        BotMessage dbMessage = new BotMessage();
        ChatEnvelope<Map<String, String>> chatEnvelope = new ChatEnvelope<>();
        String jsonMessage = JSON.toJSONString(message, SerializerFeature.WriteClassName);
        if (message instanceof AiMessage) {
            dbMessage.setRole(MessageRole.ASSISTANT.getValue());

        } else if (message instanceof UserMessage) {
            dbMessage.setRole(MessageRole.USER.getValue());
        } else if (message instanceof SystemMessage) {
            dbMessage.setRole(MessageRole.SYSTEM.getValue());
        } else if (message instanceof ToolMessage) {
            dbMessage.setRole(MessageRole.TOOL.getValue());
        }
        Map<String, String> res = new HashMap<>();
        res.put("role", dbMessage.getRole());
        res.put("content", jsonMessage);
        chatEnvelope.setType(ChatType.MESSAGE);
        chatEnvelope.setPayload(res);
        chatEnvelope.setDomain(ChatDomain.SYSTEM);
        if (dbMessage.getRole().equals(MessageRole.USER.getValue())) {
            messages.remove(messages.size() - 1);
        }
        boolean sent = sseEmitter.sendMessageNeedSave(chatEnvelope);
        if (!sent) {
            IllegalStateException e = new IllegalStateException("SSE emitter has already completed when sending needSaveMessage");
            LOG.error("sendMessageNeedSave failed, role={}, message={}, exception={}",
                    dbMessage.getRole(), e.getMessage(), e.toString(), e);
            throw e;
        }
        messages.add(res);
    }

}
