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

#include <random>
#include <catch2/catch_all.hpp>

#include "../image_pusher/TcpStreamPusher.h"
#include "../image_puller/TcpImagePuller.h"

TEST_CASE("TCPImageCommTest_2Writers", "[TCP]") {
    const size_t nframes = 128;
    const int64_t npullers = 2;
    const int64_t images_per_file = 16;

    DiffractionExperiment x(DetJF(1));
    x.Raw();
    x.PedestalG0Frames(0).NumTriggers(1).UseInternalPacketGenerator(false).IncidentEnergy_keV(12.4)
            .ImagesPerTrigger(nframes).Compression(CompressionAlgorithm::NO_COMPRESSION);

    std::mt19937 g1(1387);
    std::uniform_int_distribution<uint16_t> dist;
    std::vector<uint16_t> image1(x.GetPixelsNum() * nframes);
    for (auto &i : image1) i = dist(g1);

    std::vector<std::string> addr{
        "tcp://127.0.0.1:19001",
        "tcp://127.0.0.1:19002"
};

    std::vector<std::unique_ptr<TCPImagePuller>> puller;
    for (int i = 0; i < npullers; i++) {
        puller.push_back(std::make_unique<TCPImagePuller>(
                addr[i], 64 * 1024 * 1024)); // decoded cbor ring
    }

    TCPStreamPusher pusher(
            addr,
            64 * 1024 * 1024,
            128 * 1024, // zerocopy threshold
            8192        // sender queue
    );

    std::vector<size_t> received(npullers, 0);

    std::thread sender([&] {
        std::vector<uint8_t> serialization_buffer(16 * 1024 * 1024);
        CBORStream2Serializer serializer(serialization_buffer.data(), serialization_buffer.size());

        StartMessage start{
                .images_per_file = images_per_file,
                .write_master_file = true
        };
        EndMessage end{};

        pusher.StartDataCollection(start);

        for (int64_t i = 0; i < static_cast<int64_t>(nframes); i++) {
            DataMessage data_message;
            data_message.number = i;
            data_message.image = CompressedImage(image1.data() + i * x.GetPixelsNum(),
                                                 x.GetPixelsNum() * sizeof(uint16_t),
                                                 x.GetXPixelsNum(),
                                                 x.GetYPixelsNum(),
                                                 x.GetImageMode(),
                                                 x.GetCompressionAlgorithm());
            serializer.SerializeImage(data_message);
            REQUIRE(pusher.SendImage(serialization_buffer.data(), serializer.GetBufferSize(), i));
        }

        REQUIRE(pusher.EndDataCollection(end));
    });

    for (int w = 0; w < npullers; w++) {
        bool seen_end = false;
        while (!seen_end) {
            auto out = puller[w]->PollImage(std::chrono::seconds(10));
            REQUIRE(out.has_value());
            REQUIRE(out->cbor != nullptr);
            if (out->cbor->end_message) {
                seen_end = true;
                continue;
            }
            if (out->cbor->data_message) {
                auto n = out->cbor->data_message->number;
                REQUIRE(((n / images_per_file) % npullers) == w);
                received[w]++;
            }
        }
    }

    sender.join();

    REQUIRE(received[0] == nframes / 2);
    REQUIRE(received[1] == nframes / 2);

    for (auto &p : puller)
        p->Disconnect();
}