package tech.aiflowy.ai.entity;

import tech.aiflowy.ai.mapper.AiBotConversationMessageMapper;
import tech.aiflowy.ai.service.AiBotConversationMessageService;
import tech.aiflowy.ai.service.AiBotMessageService;
import cn.hutool.core.collection.CollectionUtil;
import cn.hutool.core.util.StrUtil;
import com.agentsflex.core.llm.functions.Function;
import com.agentsflex.core.memory.ChatMemory;
import com.agentsflex.core.message.*;
import com.alibaba.fastjson.JSON;
import com.alibaba.fastjson.serializer.SerializerFeature;
import com.mybatisflex.core.query.QueryWrapper;
import tech.aiflowy.common.satoken.util.SaTokenUtil;

import javax.annotation.Resource;
import java.math.BigInteger;
import java.util.ArrayList;
import java.util.Date;
import java.util.List;

public class AiBotMessageMemory implements ChatMemory {
    private final BigInteger botId;
    private final BigInteger accountId;
    private final String sessionId;
    private final int isExternalMsg;
    private final AiBotMessageService messageService;
    private final AiBotConversationMessageMapper aiBotConversationMessageMapper;
    private final AiBotConversationMessageService aiBotConversationService;
    public AiBotMessageMemory(BigInteger botId, BigInteger accountId, String sessionId, int isExternalMsg,
                              AiBotMessageService messageService,
                              AiBotConversationMessageMapper aiBotConversationMessageMapper,
                              AiBotConversationMessageService aiBotConversationService ) {
        this.botId = botId;
        this.accountId = accountId;
        this.sessionId = sessionId;
        this.isExternalMsg = isExternalMsg;
        this.messageService = messageService;
        this.aiBotConversationMessageMapper = aiBotConversationMessageMapper;
        this.aiBotConversationService = aiBotConversationService;
    }

    @Override
    public List<Message> getMessages() {
        List<AiBotMessage> sysAiMessages = messageService.list(QueryWrapper.create()
                .eq(AiBotMessage::getBotId, botId, true)
                .eq(AiBotMessage::getAccountId, accountId, true)
                .eq(AiBotMessage::getSessionId, sessionId, true)
                .eq(AiBotMessage::getIsExternalMsg, isExternalMsg, true)
                .orderBy(AiBotMessage::getCreated, true)
        );

        if (sysAiMessages == null || sysAiMessages.isEmpty()) {
            return null;
        }

        List<Message> messages = new ArrayList<>(sysAiMessages.size());
        for (AiBotMessage aiBotMessage : sysAiMessages) {
            Message message = aiBotMessage.toMessage();
            if (message != null) messages.add(message);
        }
        return messages;
    }

    @Override
    public void addMessage(Message message) {
        AiBotMessage aiMessage = new AiBotMessage();
        aiMessage.setCreated(new Date());
        aiMessage.setBotId(botId);
        aiMessage.setAccountId(accountId);
        aiMessage.setSessionId(sessionId);
        aiMessage.setIsExternalMsg(isExternalMsg);
        if (message instanceof AiMessage) {
            AiMessage m = (AiMessage) message;
            aiMessage.setContent(m.getFullContent());
            aiMessage.setRole("assistant");
            aiMessage.setTotalTokens(m.getTotalTokens());
            aiMessage.setPromptTokens(m.getPromptTokens());
            aiMessage.setCompletionTokens(m.getCompletionTokens());
            List<FunctionCall> calls = m.getCalls();
            if (CollectionUtil.isNotEmpty(calls)) {
                return;
            }
        } else if (message instanceof HumanMessage) {
            HumanMessage hm = (HumanMessage) message;
            aiMessage.setContent(hm.getContent());
            List<Function> functions = hm.getFunctions();
            aiMessage.setFunctions(JSON.toJSONString(functions, SerializerFeature.WriteClassName));
            aiMessage.setRole("user");
        } else if (message instanceof SystemMessage) {
            aiMessage.setRole("system");
            aiMessage.setContent(((SystemMessage) message).getContent());
        }
        if (StrUtil.isNotEmpty(aiMessage.getContent())) {
            AiBotConversationMessage aiBotConversation = aiBotConversationMessageMapper.selectOneById(aiMessage.getSessionId());
            if (aiBotConversation == null  && isExternalMsg == 1){
                AiBotConversationMessage conversation = new AiBotConversationMessage();
                conversation.setSessionId(aiMessage.getSessionId());
                conversation.setTitle(aiMessage.getContent());
                conversation.setBotId(aiMessage.getBotId());
                conversation.setCreated(new Date());
                conversation.setAccountId(SaTokenUtil.getLoginAccount().getId());
                aiBotConversationService.save(conversation);
            }
            messageService.save(aiMessage);
        }
    }

    @Override
    public Object id() {
        return botId;
    }
}