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

#include "HDF5ImageSource.h"
#include "../common/JFJochException.h"

void HDF5ImageSource::Configure(HDF5ImageLocator::Layout layout) {
    locator_.Configure(std::move(layout));
}

void HDF5ImageSource::Clear() {
    locator_.Clear();
}

HDF5ImageLocator::Location HDF5ImageSource::Resolve(int64_t global) const {
    return locator_.Resolve(global);
}

std::vector<HDF5DataSourceMessage> HDF5ImageSource::GetSourceMapping(uint64_t first_image,
                                                                     std::optional<uint64_t> image_count,
                                                                     uint64_t total_images) const {
    return locator_.GetSourceMapping(first_image, image_count, total_images);
}

CompressedImage HDF5ImageSource::ReadImageAt(std::vector<uint8_t> &buffer,
                                             const HDF5ImageLocator::Location &loc) const {
    return LoadImageDataset(buffer, *loc.file, loc.local_index);
}

CompressedImage HDF5ImageSource::LoadImageDataset(std::vector<uint8_t> &tmp, HDF5Object &file, hsize_t number) {
    std::vector<hsize_t> start = {static_cast<hsize_t>(number), 0, 0};

    const std::string dataset_name = "/entry/data/data";

    HDF5DataSet dataset(file, dataset_name);
    HDF5DataSpace dataspace(dataset);
    HDF5DataType datatype(dataset);
    HDF5Dcpl dcpl(dataset);

    if (dataspace.GetNumOfDimensions() != 3)
        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
                              "/entry/data/data dataset must be 3D");

    auto dim = dataspace.GetDimensions();

    CompressionAlgorithm algorithm = CompressionAlgorithm::NO_COMPRESSION;
    auto chunk_size = dcpl.GetChunking();

    if ((chunk_size.size() == 3) && (chunk_size[0] == 1) && (chunk_size[1] == dim[1]) && (chunk_size[2] == dim[2])) {
        dataset.ReadDirectChunk(tmp, start);
        algorithm = dcpl.GetCompression();
    } else {
        dataset.ReadVectorToU8(tmp, start, {1, dim[1], dim[2]});
        algorithm = CompressionAlgorithm::NO_COMPRESSION;
    }

    if (datatype.IsFloat())
        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
                              "Float datasets not supported at this time");

    return {
        tmp, dim[2], dim[1],
        CalcImageMode(datatype.GetElemSize(), datatype.IsFloat(), datatype.IsSigned()),
        algorithm
    };
}