#include <agent_smith/message.hpp>
#include <gtest/gtest.h>

namespace agent_smith
{
namespace
{

TEST(MessageTest, SerializesAllMessageKinds)
{
    EXPECT_EQ(
        toOpenAIJson(Message(SystemMessage{"instructions"})),
        nlohmann::json({{"role", "system"}, {"content", "instructions"}}));
    EXPECT_EQ(toOpenAIJson(Message(UserMessage{"hello"})),
              nlohmann::json({{"role", "user"}, {"content", "hello"}}));

    ToolMessage tool{
        .tool_call_id = "call_1",
        .name = "weather",
        .content = "sunny",
        .is_error = true,
    };
    EXPECT_EQ(toOpenAIJson(Message(tool)),
              nlohmann::json({{"role", "tool"},
                              {"tool_call_id", "call_1"},
                              {"name", "weather"},
                              {"content", "sunny"}}));
}

TEST(MessageTest, SerializesAssistantToolCalls)
{
    AssistantMessage assistant{
        .content = "",
        .tool_calls =
            {
                ToolCall{
                    .id = "call_1",
                    .name = "weather",
                    .arguments = {{"city", "Portland"}},
                },
            },
    };

    const nlohmann::json json = toOpenAIJson(Message(assistant));

    EXPECT_EQ(json.at("role"), "assistant");
    EXPECT_EQ(json.at("content"), "");
    ASSERT_TRUE(json.contains("tool_calls"));
    ASSERT_EQ(json.at("tool_calls").size(), 1U);
    EXPECT_EQ(json.at("tool_calls")[0],
              nlohmann::json({{"id", "call_1"},
                              {"type", "function"},
                              {"function",
                               {{"name", "weather"},
                                {"arguments", "{\"city\":\"Portland\"}"}}}}));
}

TEST(MessageTest, ParsesAssistantMessageWithToolCallsAndUsage)
{
    const nlohmann::json json{
        {"role", "assistant"},
        {"content", nullptr},
        {"tool_calls",
         {{{"id", "call_1"},
           {"type", "function"},
           {"function",
            {{"name", "weather"}, {"arguments", "{\"city\":\"Portland\"}"}}}}}},
        {"usage",
         {{"prompt_tokens", 10},
          {"completion_tokens", 4},
          {"total_tokens", 14}}},
    };

    auto parsed = assistantMessageFromOpenAIJson(json);

    ASSERT_TRUE(parsed.has_value());
    const auto* message = std::get_if<AssistantMessage>(&*parsed);
    ASSERT_NE(message, nullptr);
    EXPECT_EQ(message->content, "");
    ASSERT_EQ(message->tool_calls.size(), 1U);
    EXPECT_EQ(message->tool_calls[0].id, "call_1");
    EXPECT_EQ(message->tool_calls[0].name, "weather");
    EXPECT_EQ(message->tool_calls[0].arguments,
              nlohmann::json({{"city", "Portland"}}));
    ASSERT_TRUE(message->usage.has_value());
    EXPECT_EQ(message->usage->prompt, 10U);
    EXPECT_EQ(message->usage->completion, 4U);
    EXPECT_EQ(message->usage->total, 14U);
}

TEST(MessageTest, RejectsMalformedAssistantToolArguments)
{
    const nlohmann::json json{
        {"role", "assistant"},
        {"tool_calls",
         {{{"id", "call_1"},
           {"type", "function"},
           {"function", {{"name", "weather"}, {"arguments", "not JSON"}}}}}},
    };

    auto parsed = assistantMessageFromOpenAIJson(json);

    ASSERT_FALSE(parsed.has_value());
    const ParseError* error = parsed.error().as<ParseError>();
    ASSERT_NE(error, nullptr);
    EXPECT_EQ(error->msg, "expected tool arguments JSON object");
}

TEST(MessageTest, RejectsUnexpectedAssistantContent)
{
    const nlohmann::json json{
        {"role", "assistant"},
        {"content", {"unexpected", "object"}},
    };

    auto parsed = assistantMessageFromOpenAIJson(json);

    ASSERT_FALSE(parsed.has_value());
    EXPECT_NE(parsed.error().as<ParseError>(), nullptr);
}

} // namespace
} // namespace agent_smith