#include <string>
#include <vector>

#include <agent_smith/detail/sse_parser.hpp>
#include <gtest/gtest.h>

namespace agent_smith::internal
{
namespace
{

TEST(SseFrameParserTest, EmitsDataPayloadsAndIgnoresOtherFields)
{
    SseFrameParser parser;

    const auto payloads = parser.feed("event: completion\n"
                                      "id: 7\n"
                                      "data: {\"first\":true}\n"
                                      "\n"
                                      ": keepalive\n"
                                      "data: {\"second\":true}\n"
                                      "\n");

    EXPECT_EQ(payloads, (std::vector<std::string>{"{\"first\":true}",
                                                  "{\"second\":true}"}));
    EXPECT_FALSE(parser.finished());
}

TEST(SseFrameParserTest, RetainsIncompleteFramesAcrossChunks)
{
    SseFrameParser parser;

    EXPECT_TRUE(parser.feed("data: {\"choi").empty());
    const auto payloads = parser.feed("ces\":[1]}\n\n");

    EXPECT_EQ(payloads, (std::vector<std::string>{"{\"choices\":[1]}"}));
}

TEST(SseFrameParserTest, JoinsMultipleDataLinesWithNewlines)
{
    SseFrameParser parser;

    const auto payloads = parser.feed("data: first\n"
                                      "data: second\n"
                                      "\n");

    EXPECT_EQ(payloads, (std::vector<std::string>{"first\nsecond"}));
}

TEST(SseFrameParserTest, HandlesCarriageReturnLineEndings)
{
    SseFrameParser parser;

    const auto payloads = parser.feed("data: one\r\n\r\n");

    EXPECT_EQ(payloads, (std::vector<std::string>{"one"}));
}

TEST(SseFrameParserTest, StopsAtDoneAndDiscardsLaterPayloads)
{
    SseFrameParser parser;

    const auto payloads = parser.feed("data: before\n\n"
                                      "data: [DONE]\n\n"
                                      "data: after\n\n");

    EXPECT_EQ(payloads, (std::vector<std::string>{"before"}));
    EXPECT_TRUE(parser.finished());
    EXPECT_TRUE(parser.feed("data: ignored\n\n").empty());
}

TEST(SseFrameParserTest, ProducesIdenticalPayloadsAtEveryByteSplit)
{
    const std::string stream = "event: completion\n"
                               "data: {\"choices\":[1]}\n"
                               "\n"
                               "data: {\"usage\":{\"total_tokens\":4}}\n"
                               "\n";
    const std::vector<std::string> expected{
        "{\"choices\":[1]}",
        "{\"usage\":{\"total_tokens\":4}}",
    };

    for(std::size_t split = 0; split <= stream.size(); ++split)
    {
        SseFrameParser parser;
        std::vector<std::string> payloads =
            parser.feed(stream.substr(0, split));
        const auto suffix_payloads = parser.feed(stream.substr(split));
        payloads.insert(payloads.end(), suffix_payloads.begin(),
                        suffix_payloads.end());

        EXPECT_EQ(payloads, expected) << "split at byte " << split;
    }
}

} // namespace
} // namespace agent_smith::internal