#ifndef HUMANUS_MEMORY_BASE_H
#define HUMANUS_MEMORY_BASE_H

#include "schema.h"

namespace humanus {

struct BaseMemory {
    std::deque<Message> messages;

    // Add a message to the memory
    virtual void add_message(const Message& message) {
        messages.push_back(message);
    }

    // Add multiple messages to the memory
    void add_messages(const std::vector<Message>& messages) {
        for (const auto& message : messages) {
            add_message(message);
        }
    }
    
    // Clear all messages
    virtual void clear() {
        messages.clear();
    }

    virtual std::vector<Message> get_messages(const std::string& query = "") const {
        std::vector<Message> result;
        for (const auto& message : messages) {
            result.push_back(message);
        }
        return result;
    }

    // Convert messages to list of dicts
    json to_json_list() const {
        json memory = json::array();
        for (const auto& message : messages) {
            memory.push_back(message.to_json());
        }
        return memory;
    }
};

struct Memory : BaseMemory {
    int max_messages;

    Memory(int max_messages = 30) : max_messages(max_messages) {}

    void add_message(const Message& message) override {
        messages.push_back(message);
        while (!messages.empty() && (messages.size() > max_messages || messages.front().role == "assistant" || messages.front().role == "tool")) {
            // Ensure the first message is always a user or system message
            messages.pop_front();
        }
    }
};

}

#endif // HUMANUS_MEMORY_BASE_H