// SPDX-FileCopyrightText: 2026 Filip Leonarski, Paul Scherrer Institute <filip.leonarski@psi.ch>
// SPDX-License-Identifier: GPL-3.0-only

#include "BrokerFeed.h"

#include <base64/Base64.h>
#include <httplib.h>
#include <nlohmann/json.hpp>

#include "../common/JFJochException.h"
#include "../common/Logger.h"
#include "../frame_serialize/CBORStream2Deserializer.h"

std::vector<ServerSentEvent> ServerSentEventParser::Feed(const char *bytes, size_t size) {
    pending.append(bytes, size);
    std::vector<ServerSentEvent> events;
    size_t end;
    while ((end = pending.find("\n\n")) != std::string::npos) {
        const std::string block = pending.substr(0, end);
        pending.erase(0, end + 2);
        ServerSentEvent e;
        size_t pos = 0;
        while (pos <= block.size()) {
            size_t eol = block.find('\n', pos);
            if (eol == std::string::npos)
                eol = block.size();
            const std::string line = block.substr(pos, eol - pos);
            if (line.starts_with("event: "))
                e.event = line.substr(7);
            else if (line.starts_with("data: "))
                e.data = line.substr(6);
            pos = eol + 1;
        }
        if (!e.event.empty())
            events.push_back(std::move(e));
    }
    return events;
}

namespace {
    httplib::Headers Authorization(const std::string &token) {
        return {{"Authorization", "Bearer " + token}};
    }

    std::vector<uint8_t> DecodeBase64(const std::string &s) {
        size_t n = s.size() / 4 * 3;
        if (!s.empty() && s.back() == '=')
            n--;
        if (s.size() > 1 && s[s.size() - 2] == '=')
            n--;
        std::vector<uint8_t> out(n);
        macaron::Decode(s, out);
        return out;
    }
}

BrokerFeed::BrokerFeed(const std::string &url, std::string in_token) : token(std::move(in_token)) {
    if (url.starts_with("https://"))
        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
                              "Following a collection over https is not supported by this build - use http://");
    if (!url.starts_with("http://"))
        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
                              "The broker address must start with http://");
    const size_t path = url.find('/', 7);
    host = url.substr(0, path);
    if (host.find('@') != std::string::npos)
        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
                              "The broker address must not carry credentials");
    if (path != std::string::npos)
        prefix = url.substr(path);
    while (prefix.ends_with('/'))
        prefix.pop_back();
    thread = std::thread([this] { Run(); });
}

BrokerFeed::~BrokerFeed() {
    stop = true;
    if (thread.joinable())
        thread.join();
}

StartMessage BrokerFeed::WaitForStart() {
    std::unique_lock ul(m);
    cv.wait(ul, [&] { return start.has_value() || done; });
    if (start)
        return start.value();
    throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
                          lost.empty() ? "The broker at " + host + " sent no start message" : lost);
}

void BrokerFeed::Attach(JFJochStartMessageReader &in_reader) {
    std::unique_lock ul(m);
    reader = &in_reader;
    for (const auto &[file_number, images]: files)
        reader->FileClosed(file_number, images);
    if (ended)
        reader->Ended();
    else if (done)
        reader->Lost(lost);
}

void BrokerFeed::Run() {
    Logger logger("BrokerFeed");
    httplib::Client client(host);
    // The broker writes at least a keepalive every 10 s; three missed ones mean it is gone.
    client.set_read_timeout(std::chrono::seconds(30));
    ServerSentEventParser parser;
    int status = 0;
    std::string error;
    client.Get(prefix + "/live/events", Authorization(token),
               [&](const httplib::Response &response) {
                   status = response.status;
                   return status == 200;
               },
               [&](const char *data, size_t size) {
                   for (const auto &e: parser.Feed(data, size)) {
                       const auto j = nlohmann::json::parse(e.data);
                       std::unique_lock ul(m);
                       if (e.event == "start") {
                           auto msg = CBORStream2Deserialize(DecodeBase64(j.at("cbor").get<std::string>()));
                           if (msg && msg->start_message)
                               start = msg->start_message.value();
                           else
                               error = "the broker sent no readable start message";
                       } else if (e.event == "file") {
                           files.emplace_back(j.at("file_number").get<uint64_t>(), j.at("image_count").get<uint64_t>());
                           if (reader)
                               reader->FileClosed(files.back().first, files.back().second);
                       } else if (e.event == "end") {
                           if (j.contains("error"))
                               logger.Warning("The collection ended with an error: {}", j.at("error").get<std::string>());
                           if (!start)
                               error = "The collection at " + host + " ended before it started"
                                       + (j.contains("error") ? ": " + j.at("error").get<std::string>() : "");
                           else {
                               logger.Info("The collection ended with {} images in {} data files",
                                           j.at("images").get<uint64_t>(), files.size());
                               ended = true;
                           }
                       } else if (e.event == "superseded") {
                           error = "a new collection began before this one ended";
                       }
                       cv.notify_all();
                       if (!error.empty())
                           return false;
                   }
                   return !stop;
               });

    std::unique_lock ul(m);
    if (!ended) {
        if (!error.empty())
            lost = error;
        else if (status == 401)
            lost = "Following the collection at " + host + ": the broker refused the token in JUNGFRAUJOCH_HTTP_TOKEN";
        else if (status == 404)
            lost = "No collection to follow at " + host;
        else if (status == 409)
            lost = "Following the collection at " + host + ": another reader is following it, or its writer "
                   "does not report closed data files";
        else if (status != 200)
            lost = "No answer from the broker at " + host;
        else
            lost = "Lost the broker at " + host + " before the collection ended";
        if (reader)
            reader->Lost(lost);
    } else if (reader)
        reader->Ended();
    done = true;
    cv.notify_all();
}