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

#include <catch2/catch_all.hpp>
#include <cstdint>
#include <cstring>
#include <filesystem>
#include <fstream>
#include <hdf5.h>
#include <string>
#include <vector>

#include "../reader/JFJochXFELReader.h"
#include "../reader/JFJochXTCReader.h"

// Synthetic SACLA Cheetah and Cheetah CXI files, written with the HDF5 C API the way those pipelines
// write them: whole frames as single deflated (and byte-shuffled) chunks, so the direct-chunk path is
// what is read.
namespace {

void WriteFrames(hid_t loc, const std::string &name, hid_t type, const void *data, std::vector<hsize_t> dims,
                 bool shuffle = true) {
    const hid_t space = H5Screate_simple(static_cast<int>(dims.size()), dims.data(), nullptr);
    std::vector<hsize_t> chunk = dims;
    if (chunk.size() == 3)
        chunk[0] = 1;
    const hid_t dcpl = H5Pcreate(H5P_DATASET_CREATE);
    H5Pset_chunk(dcpl, static_cast<int>(chunk.size()), chunk.data());
    if (shuffle)
        H5Pset_shuffle(dcpl);
    H5Pset_deflate(dcpl, 4);
    const hid_t ds = H5Dcreate2(loc, name.c_str(), type, space, H5P_DEFAULT, dcpl, H5P_DEFAULT);
    H5Dwrite(ds, type, H5S_ALL, H5S_ALL, H5P_DEFAULT, data);
    H5Dclose(ds);
    H5Pclose(dcpl);
    H5Sclose(space);
}

void WriteVector(hid_t loc, const std::string &name, const std::vector<double> &v) {
    const hsize_t n = v.size();
    const hid_t space = H5Screate_simple(1, &n, nullptr);
    const hid_t ds = H5Dcreate2(loc, name.c_str(), H5T_NATIVE_DOUBLE, space, H5P_DEFAULT, H5P_DEFAULT, H5P_DEFAULT);
    H5Dwrite(ds, H5T_NATIVE_DOUBLE, H5S_ALL, H5S_ALL, H5P_DEFAULT, v.data());
    H5Dclose(ds);
    H5Sclose(space);
}

void WriteScalar(hid_t loc, const std::string &name, double v) {
    const hid_t space = H5Screate(H5S_SCALAR);
    const hid_t ds = H5Dcreate2(loc, name.c_str(), H5T_NATIVE_DOUBLE, space, H5P_DEFAULT, H5P_DEFAULT, H5P_DEFAULT);
    H5Dwrite(ds, H5T_NATIVE_DOUBLE, H5S_ALL, H5S_ALL, H5P_DEFAULT, &v);
    H5Dclose(ds);
    H5Sclose(space);
}

// A pixel of the flat image, as the reader hands it out: a 16-bit signed frame without a gain in its own
// type, anything else as int32.
int32_t PixelAt(JFJochXFELReader &reader, int64_t image, size_t col, size_t row,
                CompressedImageMode mode = CompressedImageMode::Int32) {
    auto raw = reader.GetRawImage(image);
    REQUIRE(raw);
    REQUIRE(raw->image.GetMode() == mode);
    const size_t i = row * raw->image.GetWidth() + col;
    if (mode == CompressedImageMode::Int16)
        return reinterpret_cast<const int16_t *>(raw->image_buffer.data())[i];
    return reinterpret_cast<const int32_t *>(raw->image_buffer.data())[i];
}

}  // namespace

TEST_CASE("JFJochXFELReader_SACLA", "[portable]") {
    // Two sensors of 8 rows x 4 columns stacked in a 16 x 4 frame, 50 um pixels. The first lies square
    // with its first corner on the beam; the second is turned by 90 degrees with its corner 20 pixels
    // along x (posx = -1000 um), so its fast axis runs along y and its slow axis along -x.
    const std::string name = "xfel_reader_sacla.h5";
    {
        const hid_t f = H5Fcreate(name.c_str(), H5F_ACC_TRUNC, H5P_DEFAULT, H5P_DEFAULT);
        const hid_t meta = H5Gcreate2(f, "metadata", H5P_DEFAULT, H5P_DEFAULT, H5P_DEFAULT);
        WriteVector(meta, "posx_in_um", {0, -1000});
        WriteVector(meta, "posy_in_um", {0, 0});
        WriteVector(meta, "posz_in_um", {0, 0});
        WriteVector(meta, "angle_in_rad", {0, 90});   // degrees, whatever the name says
        WriteVector(meta, "pixelsizex_in_um", {50, 50});
        WriteVector(meta, "pixelsizey_in_um", {50, 50});
        WriteScalar(meta, "distance_in_mm", 60.0);
        std::vector<uint16_t> pixelmask(16 * 4, 0);
        pixelmask[1] = 1;   // first sensor, row 0, column 1
        WriteFrames(meta, "pixelmask", H5T_NATIVE_UINT16, pixelmask.data(), {16, 4});
        H5Gclose(meta);
        for (int shot = 0; shot < 2; shot++) {
            const std::string group = "tag-" + std::to_string(100 + shot);
            const hid_t g = H5Gcreate2(f, group.c_str(), H5P_DEFAULT, H5P_DEFAULT, H5P_DEFAULT);
            std::vector<int16_t> frame(16 * 4, 10);
            frame[0] = static_cast<int16_t>(20 + shot);    // first sensor (fs 0, ss 0)
            frame[8 * 4] = -5;                             // second sensor (fs 0, ss 0): below zero
            frame[8 * 4 + 1] = 33;                         // second sensor (fs 1, ss 0)
            frame[9 * 4] = 44;                             // second sensor (fs 0, ss 1)
            frame[2 * 4] = INT16_MIN;                      // first sensor (fs 0, ss 2): no measurement
            WriteFrames(g, "data", H5T_NATIVE_INT16, frame.data(), {16, 4});
            WriteScalar(g, "photon_energy_ev", 10000.0 + 10.0 * shot);
            H5Gclose(g);
        }
        H5Fclose(f);
    }

    REQUIRE(JFJochXFELReader::CanReadPanelTable(name));
    CHECK_FALSE(JFJochXFELReader::CanReadCXI(name));
    JFJochXFELReader reader;
    REQUIRE_NOTHROW(reader.ReadPanelTable({name}));
    CHECK(reader.GetNumberOfImages() == 2);
    CHECK_FALSE(reader.CountsInPhotons());
    const auto x = reader.GetDataset()->experiment;
    CHECK(x.GetIncidentEnergy_keV() == Catch::Approx(10.005));
    CHECK(x.GetDetectorDistance_mm() == Catch::Approx(60.0));
    CHECK(x.GetPixelSize_mm() == Catch::Approx(0.05));
    // The first sensor's first pixel centre is at (0.5, 0.5) from the beam, which the flat image puts on
    // a whole column: the beam at 1.5, the pixel at column and row 2.
    CHECK(x.GetBeamX_pxl() == Catch::Approx(1.5));
    CHECK(x.GetBeamY_pxl() == Catch::Approx(1.5));
    constexpr auto int16 = CompressedImageMode::Int16;
    CHECK(PixelAt(reader, 0, 2, 2, int16) == 20);
    CHECK(PixelAt(reader, 1, 2, 2, int16) == 21);
    // Second sensor: (fs 0, ss 0) centre at x = 20 - 0.5, y = 0.5; fs runs along +y, ss along -x.
    CHECK(PixelAt(reader, 0, 21, 2, int16) == -5);
    CHECK(PixelAt(reader, 0, 21, 3, int16) == 33);
    CHECK(PixelAt(reader, 0, 20, 2, int16) == 44);
    CHECK(PixelAt(reader, 0, 0, 0, int16) == 0);   // between the panels
    // The smallest int16 is no measurement, not a count (handed out as Int16, whose error value it is).
    CHECK(PixelAt(reader, 0, 2, 4, int16) == INT16_MIN);
    // The file's mask holds the first sensor's (fs 1, ss 0), at column 3, row 2.
    const auto &mask = reader.GetDataset()->pixel_mask->GetMask();
    CHECK(mask[2 * x.GetXPixelsNum() + 3] != 0);
    CHECK(mask[2 * x.GetXPixelsNum() + 2] == 0);

    reader.ApplyGain(2.0f);
    CHECK(reader.CountsInPhotons());
    CHECK(PixelAt(reader, 0, 2, 2) == 10);
    CHECK(PixelAt(reader, 0, 21, 2) == -3);   // -2.5, rounded away from zero
    CHECK(PixelAt(reader, 0, 21, 3) == 17);   // 16.5, likewise
    reader.Close();
    std::filesystem::remove(name);
}

TEST_CASE("JFJochXFELReader_CXI", "[portable]") {
    // A single-panel Cheetah CXI file: three 6 x 6 frames of signed 32-bit counts, 75 um pixels,
    // 9 keV stated in joules, a mask_shared with one bad pixel.
    const std::string name = "xfel_reader_test.cxi";
    {
        const hid_t f = H5Fcreate(name.c_str(), H5F_ACC_TRUNC, H5P_DEFAULT, H5P_DEFAULT);
        for (const char *g: {"entry_1", "entry_1/data_1", "entry_1/instrument_1", "entry_1/instrument_1/detector_1",
                             "entry_1/instrument_1/source_1"})
            H5Gclose(H5Gcreate2(f, g, H5P_DEFAULT, H5P_DEFAULT, H5P_DEFAULT));
        std::vector<int32_t> frames(3 * 36, 5);
        frames[0] = -7;               // frame 0, pixel (0, 0): dark-subtracted below zero
        frames[36 + 7] = 70000;       // frame 1, pixel (1, 1)
        WriteFrames(f, "entry_1/data_1/data", H5T_NATIVE_INT32, frames.data(), {3, 6, 6});
        const std::string det = "entry_1/instrument_1/detector_1/";
        WriteVector(f, det + "x_pixel_size", {75e-6, 75e-6, 75e-6});
        WriteVector(f, det + "distance", {0.1, 0.1, 0.1});
        WriteVector(f, "entry_1/instrument_1/source_1/energy", std::vector<double>(3, 9000.0 * 1.602176634e-19));
        std::vector<uint16_t> mask(36, 0);
        mask[35] = 8;   // dead
        WriteFrames(f, det + "mask_shared", H5T_NATIVE_UINT16, mask.data(), {6, 6});
        H5Fclose(f);
    }

    REQUIRE(JFJochXFELReader::CanReadCXI(name));
    CHECK_FALSE(JFJochXFELReader::CanReadPanelTable(name));
    CHECK_FALSE(JFJochXFELReader::CXIRefusal(name).has_value());
    JFJochXFELReader reader;
    REQUIRE_NOTHROW(reader.ReadCXI({name}));
    CHECK(reader.GetNumberOfImages() == 3);
    const auto x = reader.GetDataset()->experiment;
    CHECK(x.GetIncidentEnergy_keV() == Catch::Approx(9.0));
    CHECK(x.GetPixelSize_mm() == Catch::Approx(0.075));
    CHECK(x.GetDetectorDistance_mm() == Catch::Approx(100.0));
    // Corner at (-3, -3): the pixel (fs, ss) lands at column and row 2 + fs, 2 + ss.
    CHECK(PixelAt(reader, 0, 2, 2) == -7);
    CHECK(PixelAt(reader, 1, 3, 3) == 70000);
    CHECK(PixelAt(reader, 2, 4, 4) == 5);
    const auto &m = reader.GetDataset()->pixel_mask->GetMask();
    CHECK(m[7 * x.GetXPixelsNum() + 7] != 0);
    CHECK(m[6 * x.GetXPixelsNum() + 6] == 0);
    reader.Close();
    std::filesystem::remove(name);
}

TEST_CASE("JFJochXFELReader_CXI_TiledRefused", "[portable]") {
    // A Cheetah CXI slab of a tiled detector is not square, and is not read as one flat panel.
    const std::string name = "xfel_reader_tiled.cxi";
    {
        const hid_t f = H5Fcreate(name.c_str(), H5F_ACC_TRUNC, H5P_DEFAULT, H5P_DEFAULT);
        for (const char *g: {"entry_1", "entry_1/data_1"})
            H5Gclose(H5Gcreate2(f, g, H5P_DEFAULT, H5P_DEFAULT, H5P_DEFAULT));
        std::vector<int16_t> frames(2 * 8 * 4, 1);
        WriteFrames(f, "entry_1/data_1/data", H5T_NATIVE_INT16, frames.data(), {2, 8, 4});
        H5Fclose(f);
    }
    CHECK_FALSE(JFJochXFELReader::CanReadCXI(name));
    const auto refusal = JFJochXFELReader::CXIRefusal(name);
    REQUIRE(refusal.has_value());
    CHECK(refusal->find("tiled detector") != std::string::npos);
    std::filesystem::remove(name);
}

TEST_CASE("JFJochXFELReader_CXI_Int16", "[portable]") {
    // A single-panel Cheetah CXI file as a Rayonix one is written: 16-bit signed frames, deflated without
    // a shuffle. The second frame's chunk is stored with its deflate skipped (filter mask bit 0), which the
    // direct-chunk decoding does not expect: HDF5 reads that one.
    const std::string name = "xfel_reader_int16.cxi";
    {
        const hid_t f = H5Fcreate(name.c_str(), H5F_ACC_TRUNC, H5P_DEFAULT, H5P_DEFAULT);
        for (const char *g: {"entry_1", "entry_1/data_1", "entry_1/instrument_1", "entry_1/instrument_1/detector_1",
                             "entry_1/instrument_1/source_1"})
            H5Gclose(H5Gcreate2(f, g, H5P_DEFAULT, H5P_DEFAULT, H5P_DEFAULT));
        std::vector<int16_t> frames(2 * 16, 3);
        frames[1] = 32767;          // frame 0, pixel (1, 0)
        frames[4 + 2] = -4;         // frame 0, pixel (2, 1)
        WriteFrames(f, "entry_1/data_1/data", H5T_NATIVE_INT16, frames.data(), {2, 4, 4}, false);
        std::vector<int16_t> raw(16, 9);
        raw[5] = 1234;              // frame 1, pixel (1, 1)
        const hid_t ds = H5Dopen2(f, "entry_1/data_1/data", H5P_DEFAULT);
        const hsize_t offset[3] = {1, 0, 0};
        REQUIRE(H5Dwrite_chunk(ds, H5P_DEFAULT, 1, offset, raw.size() * sizeof(int16_t), raw.data()) >= 0);
        H5Dclose(ds);
        const std::string det = "entry_1/instrument_1/detector_1/";
        WriteVector(f, det + "x_pixel_size", {100e-6, 100e-6});
        WriteVector(f, "entry_1/instrument_1/source_1/energy", std::vector<double>(2, 12000.0 * 1.602176634e-19));
        H5Fclose(f);
    }

    REQUIRE(JFJochXFELReader::CanReadCXI(name));
    JFJochXFELReader reader;
    REQUIRE_NOTHROW(reader.ReadCXI({name}));
    // Corner at (-2, -2): the pixel (fs, ss) lands at column and row 2 + fs, 2 + ss, in its stored type.
    constexpr auto int16 = CompressedImageMode::Int16;
    CHECK(PixelAt(reader, 0, 3, 2, int16) == 32767);
    CHECK(PixelAt(reader, 0, 4, 3, int16) == -4);
    CHECK(PixelAt(reader, 0, 5, 5, int16) == 3);
    CHECK(PixelAt(reader, 0, 1, 1, int16) == 0);   // the margin
    CHECK(PixelAt(reader, 1, 3, 3, int16) == 1234);
    CHECK(PixelAt(reader, 1, 2, 2, int16) == 9);
    CHECK(reader.GetDataset()->experiment.GetSaturationLimit() == 32767);
    reader.Close();
    std::filesystem::remove(name);
}

namespace {

// An LCLS-I XTC stream of the shape JFJochXTCReader reads: a Configure datagram, then one event (an
// L1Accept) per image, each a container holding one 16-bit camera frame (Camera::FrameV1).
void WriteXTC(const std::string &name, uint32_t width, uint32_t height, uint32_t pixel_offset,
              const std::vector<std::vector<uint16_t>> &images) {
    std::ofstream f(name, std::ios::binary);
    auto words = [&](std::initializer_list<uint32_t> w) {
        for (uint32_t v: w)
            f.write(reinterpret_cast<const char *>(&v), 4);
    };
    constexpr uint32_t ID_XTC = 1, ID_FRAME = 2, XTC_HEADER = 20;
    // clock (2 words), timestamp (2 words, the transition in bits 24-27 of the first), env, Xtc header
    words({0, 0, 0u << 24, 0, 0, 0, 0, 0, ID_XTC, XTC_HEADER});
    const uint32_t frame_extent = XTC_HEADER + 16 + 2 * width * height;
    for (const auto &image: images) {
        words({0, 0, 12u << 24, 0, 0, 0, 0, 0, ID_XTC, XTC_HEADER + frame_extent});
        words({0, 0, 7, ID_FRAME | (1u << 16), frame_extent});
        words({width, height, 16, pixel_offset});
        f.write(reinterpret_cast<const char *>(image.data()), static_cast<std::streamsize>(2 * image.size()));
    }
}

}  // namespace

TEST_CASE("JFJochXTCReader_StoredPixels", "[portable]") {
    // Two streams of 3 x 2 images. Without an offset to take off, the main loop's read hands the pixels up
    // as stored (UINT16_WIDENED), the same values the decoded int32 image holds; with one, it is taken off.
    const std::string a = "xtc_reader_a.xtc", b = "xtc_reader_b.xtc", c = "xtc_reader_c.xtc";
    WriteXTC(a, 3, 2, 0, {{1, 2, 3, 4, 5, 65535}, {7, 7, 7, 7, 7, 7}});
    WriteXTC(b, 3, 2, 0, {{9, 8, 7, 6, 5, 4}});
    WriteXTC(c, 3, 2, 10, {{11, 12, 13, 14, 15, 9}});
    REQUIRE(JFJochXTCReader::CanRead(a));

    JFJochXTCReader reader;
    REQUIRE_NOTHROW(reader.ReadFiles({a, b}));
    REQUIRE(reader.GetNumberOfImages() == 3);
    auto raw = reader.GetRawImage(2);
    REQUIRE(raw);
    CHECK(raw->image.GetCompressionAlgorithm() == CompressionAlgorithm::UINT16_WIDENED);
    std::vector<uint8_t> buffer;
    auto pixels = reinterpret_cast<const int32_t *>(raw->image.GetUncompressedPtr(buffer));
    CHECK(pixels[0] == 9);
    CHECK(pixels[5] == 4);
    raw = reader.GetRawImage(0);
    REQUIRE(raw);
    pixels = reinterpret_cast<const int32_t *>(raw->image.GetUncompressedPtr(buffer));
    CHECK(pixels[5] == 65535);

    REQUIRE_NOTHROW(reader.ReadFiles({c}));
    raw = reader.GetRawImage(0);
    REQUIRE(raw);
    CHECK(raw->image.GetCompressionAlgorithm() == CompressionAlgorithm::NO_COMPRESSION);
    CHECK(raw->image.GetMode() == CompressedImageMode::Int32);
    pixels = reinterpret_cast<const int32_t *>(raw->image_buffer.data());
    CHECK(pixels[0] == 1);
    CHECK(pixels[5] == -1);
    reader.Close();
    for (const auto &n: {a, b, c})
        std::filesystem::remove(n);
}