package tech.easyflow.ai.easyagents.memory;

import com.easyagents.core.memory.ChatMemory;
import com.easyagents.core.message.AiMessage;
import com.easyagents.core.message.Message;
import com.easyagents.core.message.SystemMessage;
import com.easyagents.core.message.UserMessage;
import tech.easyflow.core.runtime.ChatRuntimeMessage;

import java.util.ArrayList;
import java.util.List;

/**
 * 从聊天运行时消息恢复 easy-agents 历史消息。
 */
public class RuntimeChatMemory implements ChatMemory {

    private final Object id;
    private final List<Message> messages = new ArrayList<>();

    public RuntimeChatMemory(Object id, List<ChatRuntimeMessage> runtimeMessages) {
        this.id = id;
        if (runtimeMessages != null) {
            for (ChatRuntimeMessage runtimeMessage : runtimeMessages) {
                this.messages.addAll(toMessages(runtimeMessage));
            }
        }
    }

    @Override
    public List<Message> getMessages(int count) {
        if (messages.isEmpty()) {
            return null;
        }
        if (count <= 0 || messages.size() <= count) {
            return new ArrayList<>(messages);
        }
        return new ArrayList<>(messages.subList(Math.max(messages.size() - count, 0), messages.size()));
    }

    @Override
    public void addMessage(Message message) {
        if (message != null) {
            messages.add(message);
        }
    }

    @Override
    public void clear() {
        messages.clear();
    }

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

    private List<Message> toMessages(ChatRuntimeMessage runtimeMessage) {
        if (runtimeMessage == null) {
            return List.of();
        }
        String role = runtimeMessage.getRole();
        if ("assistant".equalsIgnoreCase(role)) {
            List<Message> structuredMessages = ChatRuntimeHistoryMessageMapper.toStructuredMessages(runtimeMessage);
            if (!structuredMessages.isEmpty()) {
                return structuredMessages;
            }
            if (runtimeMessage.getContentText() == null || runtimeMessage.getContentText().isBlank()) {
                return List.of();
            }
            return List.of(new AiMessage(runtimeMessage.getContentText()));
        }
        if ("system".equalsIgnoreCase(role)) {
            if (runtimeMessage.getContentText() == null || runtimeMessage.getContentText().isBlank()) {
                return List.of();
            }
            return List.of(new SystemMessage(runtimeMessage.getContentText()));
        }
        if ("tool".equalsIgnoreCase(role)) {
            List<Message> structuredMessages = ChatRuntimeHistoryMessageMapper.toStructuredMessages(runtimeMessage);
            if (!structuredMessages.isEmpty()) {
                return structuredMessages;
            }
            if (runtimeMessage.getContentText() == null || runtimeMessage.getContentText().isBlank()) {
                return List.of();
            }
            return List.of(new SystemMessage(runtimeMessage.getContentText()));
        }
        if (runtimeMessage.getContentText() == null || runtimeMessage.getContentText().isBlank()) {
            return List.of();
        }
        return List.of(new UserMessage(runtimeMessage.getContentText()));
    }
}
