#pragma once

#include <cstddef>
#include <optional>
#include <vector>

#include <mw/task.hpp>

#include "message.hpp"

namespace agent_smith
{

/// In-process conversation storage. Access from one agent at a time.
class Memory
{
public:
    /// Destroy a memory strategy through its interface.
    virtual ~Memory() = default;
    /// Record a message in conversation order.
    virtual void append(Message message) = 0;
    /// Copy history with the supplied system message prepended.
    /// Keep base_system alive until the returned task completes.
    virtual mw::Task<mw::E<std::vector<Message>>>
    snapshot(const SystemMessage& base_system) = 0;
    /// Return the retained message count, excluding the supplied system prompt.
    virtual std::size_t size() const = 0;
    /// Remove all retained messages.
    virtual void clear() = 0;
    /// Remove tool results without a matching call in the latest assistant
    /// turn. Each call ID can match only one result; order is preserved.
    static void stripOrphans(std::vector<Message>& messages);
};

/// Retains every appended message without truncation or rewriting.
class FullMemory final : public Memory
{
public:
    /// Append a message without imposing a budget.
    void append(Message message) override;
    /// Return an independent copy of the full history with the system prompt.
    mw::Task<mw::E<std::vector<Message>>>
    snapshot(const SystemMessage& base_system) override;
    /// Return the full history's message count.
    std::size_t size() const override;
    /// Discard the full history.
    void clear() override;
private:
    std::vector<Message> messages_;
};

/// Limits on retained history; every configured limit applies.
struct RollingWindowConfig
{
    /// Maximum message count; absent means unlimited.
    std::optional<std::size_t> max_messages = std::nullopt;
    /// Maximum UTF-8 byte count of serialized OpenAI messages (no separators).
    /// This conservative character estimate includes tool names and arguments.
    std::optional<std::size_t> max_chars = std::nullopt;
    /// Approximate token budget, converted to max_chars at four bytes/token.
    std::optional<std::size_t> max_tokens = std::nullopt;
};

/// Keeps recent history by dropping whole oldest messages on append.
/// Orphaned tool results are also discarded. Limits exclude the supplied
/// system prompt. A message larger than the entire budget is discarded.
class RollingWindowMemory final : public Memory
{
public:
    /// Construct a window constrained by all supplied limits (none: unlimited).
    explicit RollingWindowMemory(RollingWindowConfig config);
    /// Construct a window limited to the most recent max_messages messages.
    explicit RollingWindowMemory(std::size_t max_messages);
    /// Append a message, enforce budgets, and remove orphaned tool results.
    void append(Message message) override;
    /// Copy retained history and prepend the unbudgeted system prompt.
    mw::Task<mw::E<std::vector<Message>>>
    snapshot(const SystemMessage& base_system) override;
    /// Return the retained history's message count.
    std::size_t size() const override;
    /// Discard retained history while preserving configured limits.
    void clear() override;
private:
    RollingWindowConfig config_;
    std::vector<Message> messages_;
};

} // namespace agent_smith