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

#pragma once

#include <cstddef>
#include <cstdint>
#include <memory>
#include <new>
#include <string>
#include <utility>
#include <vector>

#include "../compression/CompressionAlgorithmEnum.h"
#include "ColorScale.h"

// std::vector value-initialises whatever it resizes into. A buffer that the very next read fills in
// full has no use for that, and on a compressed image chunk it is megabytes of memset per image. This
// allocator default-initialises instead, which for bytes is no initialisation at all. The rebind is
// not optional: without it the container rebinds to std::allocator and the zeroing comes back.
template <class T>
struct NoInitAllocator : std::allocator<T> {
    NoInitAllocator() = default;
    template <class U> NoInitAllocator(const NoInitAllocator<U> &) {}
    template <class U> struct rebind { using other = NoInitAllocator<U>; };
    template <class U> void construct(U *p) { ::new (static_cast<void *>(p)) U; }
    template <class U, class... Args> void construct(U *p, Args &&...args) {
        ::new (static_cast<void *>(p)) U(std::forward<Args>(args)...);
    }
};

// Bytes of an image as they are stored: sized by the reader and then overwritten by the read.
using RawByteBuffer = std::vector<uint8_t, NoInitAllocator<uint8_t>>;

enum class CompressedImageMode {Int8, Int16, Int32, Uint8, Uint16, Uint32, RGB, Float16, Float32, Float64};

CompressedImageMode CalcImageMode(size_t byte_depth, bool is_float, bool is_signed);

class CompressedImage {
    size_t xpixel;
    size_t ypixel;
    CompressedImageMode mode;
    CompressionAlgorithm algorithm;
    const uint8_t *data;
    size_t size; // After compression
    std::string channel;
public:
    CompressedImage();
    CompressedImage(const void* data, size_t size,
                    size_t xpixel, size_t ypixel,
                    CompressedImageMode mode,
                    CompressionAlgorithm algorithm = CompressionAlgorithm::NO_COMPRESSION,
                    std::string channel = "default");

    CompressedImage(const std::vector<rgb>& input, size_t xpixel, size_t ypixel);
    CompressedImage(const std::vector<float>& input, size_t xpixel, size_t ypixel);
    CompressedImage(const std::vector<uint8_t>& input, size_t xpixel, size_t ypixel,
                    CompressedImageMode mode = CompressedImageMode::Uint8,
                    CompressionAlgorithm algorithm = CompressionAlgorithm::NO_COMPRESSION);
    CompressedImage(const std::vector<uint16_t>& input, size_t xpixel, size_t ypixel);
    CompressedImage(const std::vector<uint32_t>& input, size_t xpixel, size_t ypixel);
    CompressedImage(const std::vector<int16_t>& input, size_t xpixel, size_t ypixel);
    CompressedImage(const std::vector<int32_t>& input, size_t xpixel, size_t ypixel);
    CompressedImage &Channel(std::string channel);

    [[nodiscard]] size_t GetUncompressedSize() const;
    [[nodiscard]] bool IsSigned() const;
    void GetUncompressed(std::vector<uint8_t> &buffer) const;
    [[nodiscard]] const uint8_t* GetUncompressedPtr(std::vector<uint8_t> &buffer) const;
    [[nodiscard]] size_t GetByteDepth() const;
    [[nodiscard]] size_t GetNumChannels() const;
    [[nodiscard]] std::string GetChannel() const;
    [[nodiscard]] size_t GetWidth() const;
    [[nodiscard]] size_t GetHeight() const;
    [[nodiscard]] CompressedImageMode GetMode() const;
    [[nodiscard]] CompressionAlgorithm GetCompressionAlgorithm() const;
    [[nodiscard]] const uint8_t *GetCompressed() const;
    [[nodiscard]] size_t GetCompressedSize() const;
};