// 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 <fstream>

#include "../preview/JFJochTIFF.h"
#include "../preview/PreviewImage.h"
#include "../writer/HDF5Objects.h"
#include "../common/RawToConvertedGeometry.h"
#include "../frame_serialize/CBORStream2Serializer.h"
#include <tiffio.h>
#include <iostream>

TEST_CASE("TIFFTest","[TIFF]") {
    std::vector<uint16_t> values(512*1024);
    CompressedImage image(values, 1024, 512);
    REQUIRE_NOTHROW(WriteTIFFToString(image));
}

TEST_CASE("TIFFTest_Write_Read","[TIFF]") {
    std::vector<uint32_t> values(512*1024);
    for (int i = 0; i < values.size(); i++)
        values[i] = (i * 17 + 2);

    CompressedImage image(values, 1024, 512);

    std::string s;
    REQUIRE_NOTHROW(s = WriteTIFFToString(image));

    std::vector<uint8_t> buffer;
    CompressedImage out = ReadTIFF(s, buffer);
    REQUIRE(out.GetMode() == CompressedImageMode::Uint32);
    REQUIRE(out.GetWidth() == 1024);
    REQUIRE(out.GetHeight() == 512);
    REQUIRE(buffer.size() == values.size() * sizeof(uint32_t));
    REQUIRE(memcmp(values.data(), buffer.data(), buffer.size()) == 0);
}

TEST_CASE("TIFFTest_Write_Read_8bit","[TIFF]") {
    std::vector<uint8_t> values(512*1024);
    for (int i = 0; i < values.size(); i++)
        values[i] = static_cast<uint8_t>(i * 17 + 2);

    CompressedImage image(values, 1024, 512, CompressedImageMode::Uint8);

    std::string s;
    REQUIRE_NOTHROW(s = WriteTIFFToString(image));

    std::vector<uint8_t> buffer;
    CompressedImage out = ReadTIFF(s, buffer);
    REQUIRE(out.GetMode() == CompressedImageMode::Uint8);
    REQUIRE(out.GetWidth() == 1024);
    REQUIRE(out.GetHeight() == 512);
    REQUIRE(buffer == values);
}

TEST_CASE("TIFFTest_Write_Read_16bit","[TIFF]") {
    std::vector<uint16_t> values(512*1024);
    for (int i = 0; i < values.size(); i++)
        values[i] = static_cast<uint16_t>(i * 17 + 2);

    CompressedImage image(values, 1024, 512);

    std::string s;
    REQUIRE_NOTHROW(s = WriteTIFFToString(image));

    std::vector<uint8_t> buffer;
    CompressedImage out = ReadTIFF(s, buffer);
    REQUIRE(out.GetMode() == CompressedImageMode::Uint16);
    REQUIRE(out.GetWidth() == 1024);
    REQUIRE(out.GetHeight() == 512);
    REQUIRE(buffer.size() == values.size() * sizeof(uint16_t));
    REQUIRE(memcmp(values.data(), buffer.data(), buffer.size()) == 0);
}

TEST_CASE("TIFFTest_File","[TIFF]") {
    std::vector<uint16_t> values(512*1024);

    for (auto &i: values)
        i = 345;

    CompressedImage image(values, 1024, 512);

    REQUIRE_NOTHROW(WriteTIFFToFile("test_image.tiff", image));
}

TEST_CASE("TIFFTest_File_signed","[TIFF]") {
    std::vector<int16_t> values(512 * 1024);

    for (int i = 0; i < values.size(); i++)
        values[i] = static_cast<int16_t>(((i % 2 == 0) ? 1 : -1) * i);

    CompressedImage image(values, 1024, 512);

    REQUIRE_NOTHROW(WriteTIFFToFile("test_image_signed.tiff", image));
}

rgb rainbowColor(float t) {
    // Ensure t is in [0,1]
    t = std::max(0.0f, std::min(1.0f, t));

    // Convert to hue (0 to 6)
    float hue = t * 6.0f;

    int phase = static_cast<int>(hue);
    float fract = hue - phase;

    uint8_t p = static_cast<uint8_t>(255 * (1.0f - fract));
    uint8_t q = static_cast<uint8_t>(255 * fract);
    uint8_t full = 255;

    switch (phase) {
        case 0: return {full, q, 0};     // Red to Yellow
        case 1: return {p, full, 0};     // Yellow to Green
        case 2: return {0, full, q};     // Green to Cyan
        case 3: return {0, p, full};     // Cyan to Blue
        case 4: return {q, 0, full};     // Blue to Magenta
        case 5: return {full, 0, p};     // Magenta to Red
        default: return {full, 0, 0};    // Fallback (red)
    }
}

TEST_CASE("TIFFTest_File_rgb","[TIFF]") {
    std::vector<rgb> values(512 * 1024);

    for (int i = 0; i < values.size(); i++) {
        float t = static_cast<float>(i % 1024) / (1024 - 1);
        values[i] = rainbowColor(t);
    }
    CompressedImage image(values, 1024, 512);

    REQUIRE_NOTHROW(WriteTIFFToFile("test_image_rgb.tiff", image));
}


TEST_CASE("TIFFTest_File_float","[TIFF]") {
    std::vector<float> values(512 * 1024);
    for (int i = 0; i < values.size(); i++)
        values[i] = i / 100.0;
    CompressedImage image(values, 1024, 512);
    REQUIRE_NOTHROW(WriteTIFFToFile("test_image_float.tiff", image));
}

// A detector with a panel table is previewed as it lies on the plane: two panels stacked in the stored
// image, placed side by side with a gap, come out side by side, the gap in the gap colour.
TEST_CASE("PreviewImage_PanelLayout", "[TIFF][PanelId]") {
    DiffractionExperiment x(DetDECTRIS(64, 32, "test", ""));
    x.BitDepthImage(16).PixelSigned(false);
    DetectorModule right = DetectorModule::AsImaged(0, 16, 64, 16);
    right.ox = 80;
    right.oy = 0;
    x.Modules(DetectorModules({DetectorModule::AsImaged(0, 0, 64, 16), right}));

    std::vector<uint16_t> pixels(64 * 32, 0);
    for (int i = 64 * 16; i < 64 * 32; i++)
        pixels[i] = 1000;   // above the foreground: the top of the colour scale

    PreviewImage preview;
    preview.Configure(x, PixelMask(x));

    std::vector<uint8_t> buffer(1 << 20);
    CBORStream2Serializer serializer(buffer.data(), buffer.size());
    DataMessage message{.number = 0, .image = CompressedImage(pixels, 64, 32)};
    serializer.SerializeImage(message);
    buffer.resize(serializer.GetBufferSize());

    PreviewImageSettings settings{.saturation_value = 10, .show_beam_center = false, .show_spots = false,
                                  .format = PreviewImageFormat::TIFF};
    const std::string tiff = preview.GenerateImage(settings, buffer);
    {
        std::ofstream f("preview_panel_layout.tiff", std::ios::binary);
        f.write(tiff.data(), tiff.size());
    }

    TIFF *t = TIFFOpen("preview_panel_layout.tiff", "r");
    REQUIRE(t);
    uint32_t width = 0, height = 0;
    TIFFGetField(t, TIFFTAG_IMAGEWIDTH, &width);
    TIFFGetField(t, TIFFTAG_IMAGELENGTH, &height);
    // Plane bounds [-0.5, 143.5) x [-0.5, 15.5): raster pixel centres 0..143 x 0..15 (PlaneRasterOf)
    CHECK(width == 144);
    CHECK(height == 16);

    std::vector<rgb> row(width);
    REQUIRE(TIFFReadScanline(t, row.data(), 8) == 1);   // plane y = 8
    TIFFClose(t);
    remove("preview_panel_layout.tiff");

    ColorScale scale;
    scale.Select(settings.scale);
    const rgb gap = scale.Apply(ColorScaleSpecial::Gap);
    auto same = [](const rgb &a, const rgb &b) { return a.r == b.r && a.g == b.g && a.b == b.b; };
    CHECK(same(row[10], scale.ApplyLUTIndex(0)));           // left panel: zero counts
    CHECK(same(row[70], gap));                         // the gap between the panels
    CHECK(same(row[100], scale.ApplyLUTIndex(scale.LUTData().size() - 1)));           // right panel: the top of the scale
}