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

#include "CompressedImage.h"

#include <utility>
#include "JFJochDecompress.h"

CompressedImage::CompressedImage()
        : data(nullptr),
          size(0),
          xpixel(0),
          ypixel(0),
          mode(CompressedImageMode::Uint8),
          algorithm(CompressionAlgorithm::NO_COMPRESSION),
          channel("default"){}

size_t CompressedImage::GetUncompressedSize() const {
    return xpixel * ypixel * GetByteDepth() * GetNumChannels();
}

bool CompressedImage::IsSigned() const {
    switch (mode) {
        case CompressedImageMode::Int8:
        case CompressedImageMode::Int16:
        case CompressedImageMode::Int32:
        case CompressedImageMode::Float16:
        case CompressedImageMode::Float32:
        case CompressedImageMode::Float64:
            return true;
        case CompressedImageMode::Uint8:
        case CompressedImageMode::Uint16:
        case CompressedImageMode::Uint32:
        case CompressedImageMode::RGB:
            return false;
        default:
            throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
                                  "Image mode not implemented");
    }
}

size_t CompressedImage::GetByteDepth() const {
    switch (mode) {
        case CompressedImageMode::Int8:
        case CompressedImageMode::Uint8:
        case CompressedImageMode::RGB:
            return 1;
        case CompressedImageMode::Int16:
        case CompressedImageMode::Uint16:
        case CompressedImageMode::Float16:
            return 2;
        case CompressedImageMode::Int32:
        case CompressedImageMode::Uint32:
        case CompressedImageMode::Float32:
            return 4;
        case CompressedImageMode::Float64:
            return 8;
        default:
            throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
                                  "Image mode not implemented");
    }
}

size_t CompressedImage::GetNumChannels() const {
    switch (mode) {
        case CompressedImageMode::Int8:
        case CompressedImageMode::Int16:
        case CompressedImageMode::Int32:
        case CompressedImageMode::Uint8:
        case CompressedImageMode::Uint16:
        case CompressedImageMode::Uint32:
        case CompressedImageMode::Float16:
        case CompressedImageMode::Float32:
        case CompressedImageMode::Float64:
            return 1;
        case CompressedImageMode::RGB:
            return 3;
        default:
            throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
                                  "Image mode not implemented");
    }
}

void CompressedImage::GetUncompressed(std::vector<uint8_t> &buffer) const {
    buffer.resize(GetUncompressedSize());
    JFJochDecompressPtr(buffer.data(),
                        algorithm, data, size,
                        xpixel * ypixel * GetNumChannels(),
                        GetByteDepth());
}

const uint8_t *CompressedImage::GetUncompressedPtr(std::vector<uint8_t> &buffer) const {
    if (algorithm == CompressionAlgorithm::NO_COMPRESSION)
        return data;
    else {
        GetUncompressed(buffer);
        return buffer.data();
    }
}

CompressedImageMode CalcImageMode(size_t byte_depth, bool is_float, bool is_signed) {
   if (is_float) {
    switch (byte_depth) {
        case 2:
            return CompressedImageMode::Float16;
        case 4:
            return CompressedImageMode::Float32;
        case 8:
            return CompressedImageMode::Float64;
        default:
            throw JFJochException(JFJochExceptionCategory::InputParameterInvalid, "Not supported");
    }
   } else if (is_signed) {
       switch (byte_depth) {
           case 1:
               return CompressedImageMode::Int8;
           case 2:
               return CompressedImageMode::Int16;
           case 4:
               return CompressedImageMode::Int32;
           default:
               throw JFJochException(JFJochExceptionCategory::InputParameterInvalid, "Not supported");
       }
   } else {
       switch (byte_depth) {
           case 1:
               return CompressedImageMode::Uint8;
           case 2:
               return CompressedImageMode::Uint16;
           case 4:
               return CompressedImageMode::Uint32;
           default:
               throw JFJochException(JFJochExceptionCategory::InputParameterInvalid, "Not supported");

       }
   }
}


CompressedImage::CompressedImage(const void *in_data, size_t in_size,
                                 size_t in_xpixel, size_t in_ypixel,
                                 CompressedImageMode in_mode,
                                 CompressionAlgorithm in_algorithm,
                                 std::string in_channel)
        : data(static_cast<const uint8_t *>(in_data)),
          size(in_size),
          xpixel(in_xpixel),
          ypixel(in_ypixel),
          mode(in_mode),
          algorithm(in_algorithm),
          channel(std::move(in_channel)) {
    if (xpixel == 0 || ypixel == 0)
        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
                              "xpixel/ypixel cannot be zero");
    if (data == nullptr) {
        if (size != 0)
            throw JFJochException(JFJochExceptionCategory::InputParameterInvalid, "For empty image, size must be 0");
    } else if (algorithm == CompressionAlgorithm::NO_COMPRESSION && GetUncompressedSize() != size)
        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
                              "Mismatch in size for uncompressed image " + std::to_string(size)
                              + " expected size "
                              + std::to_string(GetUncompressedSize()));
}

CompressedImage::CompressedImage(const std::vector<rgb> &input, size_t xpixel, size_t ypixel)
        : CompressedImage(input.data(), input.size() * sizeof(rgb),
                          xpixel, ypixel, CompressedImageMode::RGB) {}

CompressedImage::CompressedImage(const std::vector<float> &input, size_t xpixel, size_t ypixel)
        : CompressedImage(input.data(), input.size() * sizeof(float),
                          xpixel, ypixel, CompressedImageMode::Float32) {}

CompressedImage::CompressedImage(const std::vector<uint16_t> &input, size_t xpixel, size_t ypixel)
        : CompressedImage(input.data(), input.size() * sizeof(uint16_t),
                          xpixel, ypixel, CompressedImageMode::Uint16) {}

CompressedImage::CompressedImage(const std::vector<uint32_t> &input, size_t xpixel, size_t ypixel)
        : CompressedImage(input.data(), input.size() * sizeof(uint32_t),
                          xpixel, ypixel, CompressedImageMode::Uint32) {}

CompressedImage::CompressedImage(const std::vector<int16_t> &input, size_t xpixel, size_t ypixel)
        : CompressedImage(input.data(), input.size() * sizeof(int16_t),
                          xpixel, ypixel, CompressedImageMode::Int16) {}

CompressedImage::CompressedImage(const std::vector<int32_t> &input, size_t xpixel, size_t ypixel)
        : CompressedImage(input.data(), input.size() * sizeof(int32_t),
                          xpixel, ypixel, CompressedImageMode::Int32) {}

CompressedImage &CompressedImage::Channel(std::string in_channel) {
    channel = std::move(in_channel);
    return *this;
}

CompressedImage::CompressedImage(const std::vector<uint8_t> &input, size_t xpixel, size_t ypixel,
                                 CompressedImageMode mode, CompressionAlgorithm algorithm)
        : CompressedImage(input.data(), input.size(),xpixel, ypixel, mode,
                          algorithm) {}

std::string CompressedImage::GetChannel() const {
    return channel;
}

size_t CompressedImage::GetWidth() const {
    return xpixel;
}

size_t CompressedImage::GetHeight() const {
    return ypixel;
}

CompressedImageMode CompressedImage::GetMode() const {
    return mode;
}

CompressionAlgorithm CompressedImage::GetCompressionAlgorithm() const {
    return algorithm;
}

const uint8_t *CompressedImage::GetCompressed() const {
    return data;
}

size_t CompressedImage::GetCompressedSize() const {
    return size;
}