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

#include <algorithm>
#include <thread>
#include <catch2/catch_all.hpp>

#include "../preview/ZMQMetadataSocket.h"

#include "../frame_serialize/tinycbor/cbor.h"
#include "../frame_serialize/CBORStream2Deserializer.h"

TEST_CASE("ZMQMetadataSocket","[ZeroMQ]") {
    ZMQMetadataSocket metadata_socket("tcp://0.0.0.0:*");
    metadata_socket.Period(std::chrono::seconds(60));
    ZMQSocket sub_socket(ZMQSocketType::Sub);
    sub_socket.Connect(metadata_socket.GetAddress());
    sub_socket.SubscribeAll();
    sub_socket.ReceiveTimeout(std::chrono::seconds(2));

    // ensure sub socket is properly connected
    std::this_thread::sleep_for(std::chrono::seconds(2));

    StartMessage start_message{};
    DataMessage data_message_0{.number = 0};
    DataMessage data_message_1{.number = 1};
    metadata_socket.StartDataCollection(start_message);
    metadata_socket.AddDataMessage(data_message_0);
    metadata_socket.AddDataMessage(data_message_1);

    // Check if there is only start message in the queue
    ZMQMessage msg1;
    REQUIRE(sub_socket.Receive(msg1, true));
    auto msg_out = CBORStream2Deserialize(msg1.data(), msg1.size());
    REQUIRE(msg_out);
    REQUIRE(msg_out->msg_type == CBORImageType::START);

    REQUIRE(!sub_socket.Receive(msg1, true));

    EndMessage msg;
    // Metadata message should be only here after calling finalize
    metadata_socket.EndDataCollection(msg);

    REQUIRE(sub_socket.Receive(msg1, true));
    msg_out = CBORStream2Deserialize(msg1.data(), msg1.size());
    REQUIRE(msg_out);
    REQUIRE(msg_out->msg_type == CBORImageType::METADATA);
    REQUIRE(msg_out->metadata);
    REQUIRE(msg_out->metadata->images.size() == 2);
    REQUIRE(msg_out->metadata->images.at(0).number == 0);
    REQUIRE(msg_out->metadata->images.at(1).number == 1);

    REQUIRE(sub_socket.Receive(msg1, true));
    msg_out = CBORStream2Deserialize(msg1.data(), msg1.size());
    REQUIRE(msg_out);
    REQUIRE(msg_out->msg_type == CBORImageType::END);
    REQUIRE(!sub_socket.Receive(msg1, true));
}
// Image threads add at the same time; batches are cut and sent outside the lock, a batch larger than
// 20 kB per image (2D azimuthal integration) still fits, and every image arrives exactly once before END.
TEST_CASE("ZMQMetadataSocket_Concurrent","[ZeroMQ]") {
    ZMQMetadataSocket metadata_socket("tcp://0.0.0.0:*");
    metadata_socket.Period(std::chrono::milliseconds(50));
    ZMQSocket sub_socket(ZMQSocketType::Sub);
    sub_socket.Connect(metadata_socket.GetAddress());
    sub_socket.SubscribeAll();
    sub_socket.ReceiveTimeout(std::chrono::seconds(2));
    std::this_thread::sleep_for(std::chrono::seconds(2));

    constexpr int nthreads = 8;
    constexpr int nimages = 500;

    metadata_socket.StartDataCollection(StartMessage{});
    std::vector<std::thread> threads;
    for (int t = 0; t < nthreads; t++)
        threads.emplace_back([&, t] {
            for (int i = 0; i < nimages; i++) {
                DataMessage msg{.number = t * nimages + i};
                msg.az_int_profile.resize(36 * 300, 1.0f);
                msg.spots.resize(1000);
                metadata_socket.AddDataMessage(msg);
            }
        });
    for (auto &t: threads)
        t.join();
    metadata_socket.EndDataCollection(EndMessage{});

    ZMQMessage msg;
    REQUIRE(sub_socket.Receive(msg, true));
    auto msg_out = CBORStream2Deserialize(msg.data(), msg.size());
    REQUIRE(msg_out->msg_type == CBORImageType::START);

    std::vector<int64_t> numbers;
    while (true) {
        REQUIRE(sub_socket.Receive(msg, true));
        msg_out = CBORStream2Deserialize(msg.data(), msg.size());
        REQUIRE(msg_out);
        if (msg_out->msg_type == CBORImageType::END)
            break;
        REQUIRE(msg_out->msg_type == CBORImageType::METADATA);
        for (const auto &i: msg_out->metadata->images) {
            CHECK(i.az_int_profile.size() == 36 * 300);
            numbers.push_back(i.number);
        }
    }
    std::sort(numbers.begin(), numbers.end());
    REQUIRE(numbers.size() == nthreads * nimages);
    for (int i = 0; i < nthreads * nimages; i++)
        REQUIRE(numbers[i] == i);
    REQUIRE(!sub_socket.Receive(msg, true));
}

// Disabled socket: START and END are still published, no image is kept or sent
TEST_CASE("ZMQMetadataSocket_Disabled","[ZeroMQ]") {
    ZMQMetadataSocket metadata_socket("tcp://0.0.0.0:*");
    metadata_socket.Period(std::nullopt);
    ZMQSocket sub_socket(ZMQSocketType::Sub);
    sub_socket.Connect(metadata_socket.GetAddress());
    sub_socket.SubscribeAll();
    sub_socket.ReceiveTimeout(std::chrono::seconds(2));
    std::this_thread::sleep_for(std::chrono::seconds(2));

    metadata_socket.StartDataCollection(StartMessage{});
    for (int i = 0; i < 100; i++)
        metadata_socket.AddDataMessage(DataMessage{.number = i});
    metadata_socket.EndDataCollection(EndMessage{});

    ZMQMessage msg;
    REQUIRE(sub_socket.Receive(msg, true));
    REQUIRE(CBORStream2Deserialize(msg.data(), msg.size())->msg_type == CBORImageType::START);
    REQUIRE(sub_socket.Receive(msg, true));
    REQUIRE(CBORStream2Deserialize(msg.data(), msg.size())->msg_type == CBORImageType::END);
    REQUIRE(!sub_socket.Receive(msg, true));
}