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

#pragma once

// CUDA only: the engines that read a preprocessed image on the device all compile as .cu. The CPU
// engines read ImagePreprocessorBuffer directly, which is int32 whatever the source was.
#include <cuda_runtime.h>

#include <cstdint>
#include <climits>

#include "ImagePreprocessorBuffer.h"

// How a preprocessed image is stored, and how every engine reads it.
//
// The preprocessed image is normally int32 per pixel, in the convention the whole pipeline shares:
// INT32_MIN is masked or bad, INT32_MAX is saturated, anything else is a real value. On a detector
// reading out 16 bits - which is what a fast acquisition uses, and what jfjoch writes for DECTRIS
// by default - that spends four bytes a pixel to carry two, and the per-pixel passes over a full
// frame are the largest single cost of the image loop.
//
// So a 16-bit source keeps its width. Two codes at the top of the range carry the two special
// states, and they are unambiguous because a detector's saturation limit bounds what a real value
// can be:
//
//   0xFFFF            masked, or the source's own bad-pixel marker
//   the saturation    a pixel at or above the saturation limit. 0xFFFE where the limit leaves room
//   code              for it - a 16-bit EIGER declares a count-rate limit of a few thousand, so
//                     there is room to spare - and 0xFFFF where the limit is the full range, in
//                     which case the "is error" test has already claimed 0xFFFF and no pixel can
//                     be saturated at all, so 0xFFFE stays a real value.
//
// Either way a real value is strictly below the saturation limit and so below both codes. Nothing
// is clipped and nothing is lost.
//
// Engines do not learn a second convention. PixelView widens on load, so a masked pixel still reads
// as INT32_MIN and every existing `v != INT32_MIN && v != INT32_MAX` test keeps its meaning.
namespace preprocessed_pixel {
    // Masked or bad. Always this value: a 16-bit source uses it as its own bad-pixel marker.
    constexpr uint16_t NARROW_BAD = UINT16_MAX;

    // Which code means "saturated", given the saturation limit this image was preprocessed with.
    // A real value is strictly below the limit, so the limit may be as high as UINT16_MAX - 1 and
    // still leave UINT16_MAX - 1 free to be the code. Only a limit of the whole range falls back to
    // NARROW_BAD, and that is exactly the case where no pixel can be saturated - the "is error" test
    // claims UINT16_MAX first - so the code is then unreachable rather than ambiguous.
    constexpr uint16_t NarrowSaturatedCode(int64_t saturation_limit) {
        return saturation_limit <= UINT16_MAX - 1 ? static_cast<uint16_t>(UINT16_MAX - 1) : NARROW_BAD;
    }
}

struct PixelView {
    const int32_t *wide = nullptr;
    const uint16_t *narrow = nullptr;   // set instead of `wide` when the image kept its 16-bit width
    uint16_t sat_code = preprocessed_pixel::NARROW_BAD;

    __host__ __device__ __forceinline__ int32_t Widen(uint16_t v) const {
        // Order matters: where sat_code is NARROW_BAD the first test claims it, which is right -
        // that is the case where nothing can be saturated.
        if (v == preprocessed_pixel::NARROW_BAD) return INT32_MIN;
        if (v == sat_code) return INT32_MAX;
        return static_cast<int32_t>(v);
    }

    // One pixel. The branch is on a pointer that is the same for every thread of every block, so it
    // costs a predicted branch on kernels whose time is entirely the loads it selects between.
    __host__ __device__ __forceinline__ int32_t operator[](size_t i) const {
        return narrow ? Widen(narrow[i]) : wide[i];
    }

    // Pixels 4q..4q+3 in one transaction - 16 bytes wide or 8, whichever the image is. Callers that
    // use this must keep to the same quad indexing for both forms, which they do: the layout is the
    // same image either way, only narrower.
    __host__ __device__ __forceinline__ void Load4(size_t q, int32_t out[4]) const {
        if (narrow) {
            const ushort4 v = reinterpret_cast<const ushort4 *>(narrow)[q];
            out[0] = Widen(v.x); out[1] = Widen(v.y); out[2] = Widen(v.z); out[3] = Widen(v.w);
        } else {
            const int4 v = reinterpret_cast<const int4 *>(wide)[q];
            out[0] = v.x; out[1] = v.y; out[2] = v.z; out[3] = v.w;
        }
    }
};

// The only way an engine should get a view of a preprocessed image: it reads the width off the
// buffer, so the pointer and the width cannot be taken from different places and disagree.
inline PixelView ViewOf(const ImagePreprocessorBuffer &image) {
    return image.IsNarrow()
           ? PixelView{nullptr, image.getGPUBufferNarrow(), image.NarrowSatCode()}
           : PixelView{image.getGPUBuffer(), nullptr, preprocessed_pixel::NARROW_BAD};
}