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

#include <catch2/catch_all.hpp>

#include "../common/ImageBuffer.h"

TEST_CASE("ImageBuffer_Location size") {
    ImageBuffer buf(2*1024*1024);
    REQUIRE_THROWS(buf.StartMeasurement(4 * 1024 * 1024));
    REQUIRE_NOTHROW(buf.StartMeasurement(1024 * 1024));
    REQUIRE(buf.GetStatus().total_slots == 2);
    REQUIRE(buf.GetStatus().available_slots == 2);
}

TEST_CASE("ImageBuffer") {
    DiffractionExperiment experiment(DetJF(18,3,8,36));

    ImageBuffer buf_ctrl(4*experiment.GetImageBufferLocationSize());
    buf_ctrl.StartMeasurement(experiment);

    REQUIRE(buf_ctrl.GetAvailSlots() == 4);

    ZeroCopyReturnValue *ret0 = buf_ctrl.GetImageSlot();
    REQUIRE(ret0 != nullptr);
    REQUIRE(buf_ctrl.GetAvailSlots() == 3);

    ZeroCopyReturnValue *ret1 = buf_ctrl.GetImageSlot();
    REQUIRE(ret1 != nullptr);

    ZeroCopyReturnValue *ret2 = buf_ctrl.GetImageSlot();
    REQUIRE(ret2 != nullptr);

    ZeroCopyReturnValue *ret3 = buf_ctrl.GetImageSlot();
    REQUIRE(ret3 != nullptr);

    REQUIRE(buf_ctrl.GetAvailSlots() == 0);

    ZeroCopyReturnValue *ret4 = buf_ctrl.GetImageSlot();
    REQUIRE(ret4 == nullptr);

    ret0->release();
    ret1->release();
    ret2->release();
    ret3->release();

    REQUIRE(buf_ctrl.GetAvailSlots() == 4);

    REQUIRE(buf_ctrl.CheckIfBufferReturned(std::chrono::microseconds(1)));
}

TEST_CASE("ImageBuffer_Preview_Access") {
    DiffractionExperiment experiment(DetJF(18,3,8,36));

    ImageBuffer buf_ctrl(4*experiment.GetImageBufferLocationSize());
    buf_ctrl.StartMeasurement(experiment);

    REQUIRE(buf_ctrl.GetAvailSlots() == 4);

    ZeroCopyReturnValue *ret0 = buf_ctrl.GetImageSlot();
    REQUIRE(ret0 != nullptr);
    ret0->SetImageNumber(0);
    ret0->release();

    ZeroCopyReturnValue *ret1 = buf_ctrl.GetImageSlot();
    REQUIRE(ret1 != nullptr);
    ret1->SetImageNumber(1);
    ret1->release();

    ZeroCopyReturnValue *ret2 = buf_ctrl.GetImageSlot();
    REQUIRE(ret2 != nullptr);
    ret2->SetImageNumber(2);
    ret2->SetImageSize(3456);
    ret2->SetIndexed(true);
    ret2->release();

    ZeroCopyReturnValue *ret3 = buf_ctrl.GetImageSlot();
    REQUIRE(ret3 != nullptr);
    ret3->SetImageNumber(15);
    ret3->SetImageSize(1234);
    ret3->SetIndexed(false);
    ret3->release();

    REQUIRE(buf_ctrl.GetAvailSlots() == 4);

    auto v = buf_ctrl.GetStatus().images_in_the_buffer;
    REQUIRE(v.size() == 4);
    REQUIRE(v[0] == 0);
    REQUIRE(v[3] == 15);

    std::vector<uint8_t> image_2, image_2_2, image_12, image_max, image_max_indexed;
    REQUIRE(buf_ctrl.GetImage(image_2, 2));
    REQUIRE(image_2.size() == 3456);

    REQUIRE(!buf_ctrl.GetImage(image_12, 12));

    REQUIRE(buf_ctrl.Finalize(std::chrono::microseconds(1)));

    // After finalize was run, there should be no preview
    REQUIRE(!buf_ctrl.GetImage(image_2_2, 2));
}


TEST_CASE("ImageBuffer_Counter") {
    DiffractionExperiment experiment(DetJF(18,3,8,36));

    ImageBuffer buf_ctrl(4*experiment.GetImageBufferLocationSize());
    auto cnt_1 = buf_ctrl.GetStatus().current_counter;

    buf_ctrl.StartMeasurement(experiment);
    auto cnt_2 = buf_ctrl.GetStatus().current_counter;

    ZeroCopyReturnValue *ret0 = buf_ctrl.GetImageSlot();
    REQUIRE(ret0 != nullptr);
    ret0->SetImageNumber(0);
    ret0->release();

    auto cnt_3 = buf_ctrl.GetStatus().current_counter;

    ZeroCopyReturnValue *ret1 = buf_ctrl.GetImageSlot();
    REQUIRE(ret1 != nullptr);
    ret1->SetImageNumber(1);
    ret1->release();

    auto cnt_4 = buf_ctrl.GetStatus().current_counter;

    REQUIRE(buf_ctrl.Finalize(std::chrono::microseconds(1)));
    auto cnt_5 = buf_ctrl.GetStatus().current_counter;

    CHECK(cnt_1 != cnt_2);
    CHECK(cnt_1 != cnt_3);
    CHECK(cnt_1 != cnt_4);
    CHECK(cnt_1 != cnt_5);
    CHECK(cnt_2 != cnt_3);
    CHECK(cnt_2 != cnt_4);
    CHECK(cnt_2 != cnt_5);
    CHECK(cnt_3 != cnt_4);
    CHECK(cnt_3 != cnt_5);
    CHECK(cnt_4 != cnt_5);
}




TEST_CASE("ImageBuffer_Preview_Access_GetImage_MaxImage") {
    DiffractionExperiment experiment(DetJF(18,3,8,36));

    ImageBuffer buf_ctrl(4*experiment.GetImageBufferLocationSize());
    buf_ctrl.StartMeasurement(experiment);

    REQUIRE(buf_ctrl.GetAvailSlots() == 4);

    ZeroCopyReturnValue *ret0 = buf_ctrl.GetImageSlot();
    REQUIRE(ret0 != nullptr);
    ret0->SetImageNumber(0);
    ret0->release();

    ZeroCopyReturnValue *ret1 = buf_ctrl.GetImageSlot();
    REQUIRE(ret1 != nullptr);
    ret1->SetImageNumber(1);
    ret1->release();

    ZeroCopyReturnValue *ret2 = buf_ctrl.GetImageSlot();
    REQUIRE(ret2 != nullptr);
    ret2->SetImageNumber(2);
    ret2->SetImageSize(3456);
    ret2->SetIndexed(true);
    ret2->release();

    ZeroCopyReturnValue *ret3 = buf_ctrl.GetImageSlot();
    REQUIRE(ret3 != nullptr);
    ret3->SetImageNumber(15);
    ret3->SetImageSize(1234);
    ret3->SetIndexed(false);
    ret3->release();

    std::vector<uint8_t> image_max, image_max_indexed;

    REQUIRE(buf_ctrl.GetImage(image_max, ImageBuffer::MaxImage));
    REQUIRE(image_max.size() == 1234);

    REQUIRE(buf_ctrl.GetImage(image_max_indexed, ImageBuffer::MaxIndexedImage));
    REQUIRE(image_max_indexed.size() == 3456);

    REQUIRE(buf_ctrl.Finalize(std::chrono::microseconds(1)));
}

TEST_CASE("ImageBuffer_Restart") {
    DiffractionExperiment experiment(DetJF(18,3,8,36));

    ImageBuffer buf_ctrl(4*experiment.GetImageBufferLocationSize());
    buf_ctrl.StartMeasurement(experiment);

    ZeroCopyReturnValue *ret0 = buf_ctrl.GetImageSlot();
    REQUIRE(ret0 != nullptr);
    ret0->SetImageNumber(0);
    ret0->release();

    REQUIRE(buf_ctrl.GetStatus().images_in_the_buffer.size() == 1);

    // Restart measurement
    REQUIRE_NOTHROW(buf_ctrl.StartMeasurement(experiment));

    REQUIRE(buf_ctrl.GetAvailSlots() == 4);
    REQUIRE(buf_ctrl.GetStatus().images_in_the_buffer.empty());


    ret0 = buf_ctrl.GetImageSlot();
    REQUIRE(ret0 != nullptr);
    ret0->SetImageNumber(0);
    ret0->release();

    REQUIRE_NOTHROW(buf_ctrl.StartMeasurement(experiment));
}

TEST_CASE("ImageBuffer_BuffersNotReturned") {
    DiffractionExperiment experiment(DetJF(18,3,8,36));

    ImageBuffer buf_ctrl(4*experiment.GetImageBufferLocationSize());
    buf_ctrl.StartMeasurement(experiment);

    REQUIRE(buf_ctrl.GetAvailSlots() == 4);

    ZeroCopyReturnValue *ret0 = buf_ctrl.GetImageSlot();
    REQUIRE(ret0 != nullptr);
    REQUIRE(buf_ctrl.GetAvailSlots() == 3);
    REQUIRE(!buf_ctrl.CheckIfBufferReturned(std::chrono::microseconds(1)));
}