#include <limits>

#include <agent_smith/memory.hpp>
#include <agent_smith/sync_wait.hpp>
#include <gtest/gtest.h>

namespace agent_smith
{
namespace
{

std::vector<Message> snapshot(Memory& memory)
{
    const SystemMessage system{"system"};
    auto result = syncWait(memory.snapshot(system));
    EXPECT_TRUE(result);
    return result ? std::move(*result) : std::vector<Message>{};
}

AssistantMessage calls()
{
    return {.content = "calling",
            .tool_calls = {{"a", "weather", {{"city", "Tokyo"}}},
                           {"b", "weather", {{"city", "Paris"}}}}};
}

TEST(FullMemoryTest, PreservesHistoryAndCopiesSnapshots)
{
    FullMemory memory;
    EXPECT_EQ(snapshot(memory).size(), 1U);
    memory.append(UserMessage{"hi"});
    memory.append(calls());
    memory.append(ToolMessage{"a", "weather", "sunny"});
    memory.append(ToolMessage{"b", "weather", "rainy"});
    memory.append(AssistantMessage{.content = "done"});
    ASSERT_EQ(memory.size(), 5U);
    auto first = snapshot(memory);
    ASSERT_EQ(first.size(), 6U);
    EXPECT_EQ(std::get<SystemMessage>(first[0]).content, "system");
    EXPECT_EQ(std::get<AssistantMessage>(first[2]).tool_calls.size(), 2U);
    std::get<UserMessage>(first[1]).content = "changed";
    EXPECT_EQ(std::get<UserMessage>(snapshot(memory)[1]).content, "hi");
    memory.clear();
    EXPECT_EQ(memory.size(), 0U);
    EXPECT_EQ(first.size(), 6U);
}

TEST(RollingWindowMemoryTest, RetainsNewestAndKeepsSystemOutsideBudget)
{
    RollingWindowMemory memory(2);
    memory.append(UserMessage{"old"});
    memory.append(AssistantMessage{.content = "answer"});
    memory.append(UserMessage{"new"});
    auto result = snapshot(memory);
    ASSERT_EQ(result.size(), 3U);
    EXPECT_EQ(std::get<AssistantMessage>(result[1]).content, "answer");
    EXPECT_EQ(std::get<UserMessage>(result[2]).content, "new");
    memory.clear();
    EXPECT_EQ(memory.size(), 0U);
    memory.append(UserMessage{"again"});
    EXPECT_EQ(memory.size(), 1U);
}

TEST(RollingWindowMemoryTest, AppliesCharacterAndTokenBudgets)
{
    const Message message = UserMessage{"hello"};
    const auto bytes = toOpenAIJson(message).dump().size();
    RollingWindowMemory exact({.max_chars = bytes});
    exact.append(message);
    EXPECT_EQ(exact.size(), 1U);
    exact.append(message);
    EXPECT_EQ(exact.size(), 1U);
    RollingWindowMemory too_small({.max_chars = bytes - 1});
    too_small.append(message);
    EXPECT_EQ(too_small.size(), 0U);
    RollingWindowMemory tokens({.max_tokens = (bytes + 3) / 4});
    tokens.append(message);
    EXPECT_EQ(tokens.size(), 1U);
    tokens.append(message);
    EXPECT_EQ(tokens.size(), 1U);
    RollingWindowMemory both({.max_messages = 0, .max_tokens = 100});
    both.append(message);
    EXPECT_EQ(snapshot(both).size(), 1U);
    RollingWindowMemory huge({.max_tokens =
        std::numeric_limits<std::size_t>::max()});
    huge.append(message);
    EXPECT_EQ(huge.size(), 1U);
}

TEST(RollingWindowMemoryTest, CountsToolArgumentsAndDropsEntireGroup)
{
    RollingWindowMemory memory(3);
    memory.append(calls());
    memory.append(ToolMessage{"a", "weather", "sunny"});
    memory.append(ToolMessage{"b", "weather", "rainy"});
    EXPECT_EQ(memory.size(), 3U);
    memory.append(UserMessage{"next"});
    ASSERT_EQ(memory.size(), 1U);
    EXPECT_EQ(std::get<UserMessage>(snapshot(memory)[1]).content, "next");
    RollingWindowMemory small({.max_chars = 100});
    auto large = calls();
    large.tool_calls[0].arguments = {{"data", std::string(1000, 'x')}};
    small.append(large);
    small.append(ToolMessage{"a", "weather", "result"});
    EXPECT_EQ(small.size(), 0U);
}

TEST(MemoryTest, RemovesOrphansDuplicatesAndCallsFromPreviousAssistant)
{
    std::vector<Message> messages{
        ToolMessage{"a", "weather", "orphan"}, calls(),
        ToolMessage{"b", "weather", "second first"},
        ToolMessage{"a", "weather", "first second"},
        ToolMessage{"a", "weather", "duplicate"},
        AssistantMessage{.content = "done"},
        ToolMessage{"b", "weather", "stale"}};
    Memory::stripOrphans(messages);
    ASSERT_EQ(messages.size(), 4U);
    EXPECT_EQ(std::get<ToolMessage>(messages[1]).tool_call_id, "b");
    EXPECT_EQ(std::get<ToolMessage>(messages[2]).tool_call_id, "a");
    Memory::stripOrphans(messages);
    EXPECT_EQ(messages.size(), 4U);
}

TEST(RollingWindowMemoryTest, CombinedLimitsAndOversizeMessages)
{
    const Message message = UserMessage{"hello"};
    const auto bytes = toOpenAIJson(message).dump().size();
    RollingWindowMemory memory({.max_messages = 3, .max_chars = bytes * 2,
                                .max_tokens = 1000});
    for(int index = 0; index < 4; ++index)
    {
        memory.append(message);
    }
    EXPECT_EQ(memory.size(), 2U);
    memory.append(UserMessage{std::string(1000, 'x')});
    EXPECT_EQ(memory.size(), 0U);
    memory.append(message);
    EXPECT_EQ(memory.size(), 1U);
    RollingWindowMemory zero({.max_tokens = 0});
    zero.append(UserMessage{""});
    EXPECT_EQ(zero.size(), 0U);
}

TEST(MemoryTest, ChangingSystemPromptDoesNotAccumulateHistory)
{
    FullMemory full;
    RollingWindowMemory rolling(4);
    for(Memory* memory : std::vector<Memory*>{&full, &rolling})
    {
        memory->append(SystemMessage{"stored instruction"});
        memory->append(UserMessage{"你好"});
        const SystemMessage first{"first"};
        const SystemMessage second{"second"};
        auto original = syncWait(memory->snapshot(first));
        auto updated = syncWait(memory->snapshot(second));
        ASSERT_TRUE(original);
        ASSERT_TRUE(updated);
        ASSERT_EQ(updated->size(), 3U);
        EXPECT_EQ(std::get<SystemMessage>((*original)[0]).content, "first");
        EXPECT_EQ(std::get<SystemMessage>((*updated)[0]).content, "second");
        EXPECT_EQ(std::get<SystemMessage>((*updated)[1]).content,
                  "stored instruction");
        EXPECT_EQ(memory->size(), 2U);
    }
}

} // namespace
} // namespace agent_smith