#ifndef HUMANUS_MEMORY_BASE_H
#define HUMANUS_MEMORY_BASE_H

#include "schema.h"

namespace humanus {

struct BaseMemory {
    std::vector<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
    void clear() {
        messages.clear();
    }

    virtual std::vector<Message> get_messages() const {
        return messages;
    }

    // Get the last n messages
    virtual std::vector<Message> get_recent_messages(int n) const {
        n = std::min(n, static_cast<int>(messages.size()));
        return std::vector<Message>(messages.end() - n, messages.end());
    }

    // 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 = 100) : max_messages(max_messages) {}

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

}

#endif // HUMANUS_MEMORY_BASE_H