#include "agent_smith/detail/sse_parser.hpp"

#include <optional>
#include <string_view>
#include <utility>

namespace agent_smith::internal
{
namespace
{

struct FrameBoundary
{
    std::size_t position;
    std::size_t length;
};

std::optional<FrameBoundary> findFrameBoundary(std::string_view bytes)
{
    const std::size_t lf_boundary = bytes.find("\n\n");
    const std::size_t crlf_boundary = bytes.find("\r\n\r\n");

    if(lf_boundary == std::string_view::npos &&
       crlf_boundary == std::string_view::npos)
    {
        return std::nullopt;
    }
    if(crlf_boundary == std::string_view::npos ||
       (lf_boundary != std::string_view::npos && lf_boundary < crlf_boundary))
    {
        return FrameBoundary{
            .position = lf_boundary,
            .length = 2,
        };
    }
    return FrameBoundary{
        .position = crlf_boundary,
        .length = 4,
    };
}

std::optional<std::string> dataPayload(std::string_view frame)
{
    std::string payload;
    bool has_data = false;
    std::size_t line_start = 0;

    while(line_start <= frame.size())
    {
        const std::size_t line_end = frame.find('\n', line_start);
        const std::size_t length = line_end == std::string_view::npos
                                       ? frame.size() - line_start
                                       : line_end - line_start;
        std::string_view line = frame.substr(line_start, length);
        if(!line.empty() && line.back() == '\r')
        {
            line.remove_suffix(1);
        }

        if(line.starts_with("data:"))
        {
            std::string_view value = line.substr(5);
            if(value.starts_with(' '))
            {
                value.remove_prefix(1);
            }
            if(has_data)
            {
                payload.push_back('\n');
            }
            payload.append(value);
            has_data = true;
        }

        if(line_end == std::string_view::npos)
        {
            break;
        }
        line_start = line_end + 1;
    }

    if(!has_data)
    {
        return std::nullopt;
    }
    return payload;
}

} // namespace

std::vector<std::string> SseFrameParser::feed(std::string_view bytes)
{
    if(finished_)
    {
        return {};
    }

    buffer_.append(bytes);
    std::vector<std::string> payloads;

    while(true)
    {
        const auto boundary = findFrameBoundary(buffer_);
        if(!boundary)
        {
            break;
        }

        std::string frame = buffer_.substr(0, boundary->position);
        buffer_.erase(0, boundary->position + boundary->length);
        auto payload = dataPayload(frame);
        if(!payload)
        {
            continue;
        }
        if(*payload == "[DONE]")
        {
            finished_ = true;
            buffer_.clear();
            break;
        }
        payloads.push_back(std::move(*payload));
    }

    return payloads;
}

bool SseFrameParser::finished() const noexcept
{
    return finished_;
}

} // namespace agent_smith::internal