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

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

namespace {
    // Coalesce consecutive single-image mappings into one contiguous range when the source and
    // virtual images stay contiguous in the same file/dataset.
    void AppendOrExtendSourceMapping(std::vector<HDF5DataSourceMessage> &ret,
                                     const std::string &filename,
                                     const std::string &dataset,
                                     uint64_t source_first_image,
                                     uint64_t virtual_first_image,
                                     uint64_t image_count) {
        if (image_count == 0)
            return;

        if (!ret.empty()) {
            auto &last = ret.back();
            if (last.filename == filename
                && last.dataset == dataset
                && last.source_first_image + last.image_count == source_first_image
                && last.virtual_first_image + last.image_count == virtual_first_image) {
                last.image_count += image_count;
                return;
            }
        }

        ret.push_back(HDF5DataSourceMessage{
            .filename = filename,
            .dataset = dataset,
            .source_first_image = source_first_image,
            .virtual_first_image = virtual_first_image,
            .image_count = image_count
        });
    }
}

void HDF5ImageLocator::Configure(Layout layout) {
    file_cache_.clear();
    layout_ = std::move(layout);
}

void HDF5ImageLocator::Clear() {
    file_cache_.clear();
    layout_ = Layout{};
}

std::shared_ptr<HDF5ReadOnlyFile> HDF5ImageLocator::OpenCached(const std::string &path) const {
    auto it = file_cache_.find(path);
    if (it != file_cache_.end())
        return it->second;
    auto file = std::make_shared<HDF5ReadOnlyFile>(path);
    file_cache_[path] = file;
    return file;
}

HDF5ImageLocator::Location HDF5ImageLocator::Resolve(int64_t global_image) const {
    if (global_image < 0)
        throw JFJochException(JFJochExceptionCategory::HDF5, "Image out of bounds");

    if (layout_.format == FileWriterFormat::NXmxLegacy) {
        const uint32_t file_id = global_image / layout_.images_per_file;
        const uint32_t local_index = global_image % layout_.images_per_file;
        return {OpenCached(layout_.legacy_files.at(file_id)), local_index};
    }

    if (layout_.format == FileWriterFormat::NXmxVDS
        && layout_.data_layout == HDF5DataSetLayout::VIRTUAL) {
        const auto image = static_cast<hsize_t>(global_image);
        for (const auto &mapping: layout_.vds_mappings) {
            if (!mapping.ContainsVirtualImage(image))
                continue;
            return {OpenCached(mapping.filename), static_cast<uint32_t>(mapping.SourceImage(image))};
        }
        throw JFJochException(JFJochExceptionCategory::HDF5,
                              "Image not covered by /entry/data/data VDS mappings");
    }

    // Contiguous / integrated: pixels live in the master file at the global index.
    if (!layout_.master_file)
        throw JFJochException(JFJochExceptionCategory::HDF5, "Master file not loaded");
    return {layout_.master_file, static_cast<uint32_t>(global_image)};
}

std::vector<HDF5DataSourceMessage> HDF5ImageLocator::GetSourceMapping(uint64_t first_image,
                                                                      std::optional<uint64_t> image_count,
                                                                      uint64_t total_images) const {
    if (first_image > total_images)
        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
                              "First image outside dataset range");

    const uint64_t requested_count = image_count.value_or(total_images - first_image);
    if (first_image + requested_count > total_images)
        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
                              "Requested image range outside dataset range");

    std::vector<HDF5DataSourceMessage> ret;
    if (requested_count == 0)
        return ret;

    // Integrated / contiguous source: link directly to the original master file.
    if (layout_.format == FileWriterFormat::NXmxVDS && layout_.data_layout != HDF5DataSetLayout::VIRTUAL) {
        AppendOrExtendSourceMapping(ret, layout_.master_filename, "/entry/data/data",
                                    first_image, 0, requested_count);
        return ret;
    }

    // VDS source: expand VDS mappings to original source files, not to the VDS master.
    if (layout_.format == FileWriterFormat::NXmxVDS && layout_.data_layout == HDF5DataSetLayout::VIRTUAL) {
        for (uint64_t local_image = 0; local_image < requested_count; ++local_image) {
            const hsize_t virtual_image = first_image + local_image;

            bool found = false;
            for (const auto &mapping: layout_.vds_mappings) {
                if (!mapping.ContainsVirtualImage(virtual_image))
                    continue;

                const uint64_t source_image = mapping.SourceImage(virtual_image);
                const std::string dataset = mapping.dataset.empty() ? "/entry/data/data" : mapping.dataset;

                AppendOrExtendSourceMapping(ret, mapping.filename, dataset, source_image, local_image, 1);
                found = true;
                break;
            }

            if (!found)
                throw JFJochException(JFJochExceptionCategory::HDF5,
                                      "Image not covered by /entry/data/data VDS mappings");
        }

        return ret;
    }

    // Legacy source: link directly to the linked data files.
    if (layout_.format == FileWriterFormat::NXmxLegacy) {
        if (layout_.images_per_file == 0)
            throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
                                  "Cannot generate HDF5 source mapping: images_per_file is zero");

        for (uint64_t local_image = 0; local_image < requested_count; ++local_image) {
            const uint64_t source_global_image = first_image + local_image;
            const uint64_t file_id = source_global_image / layout_.images_per_file;
            const uint64_t source_image = source_global_image % layout_.images_per_file;

            if (file_id >= layout_.legacy_files.size())
                throw JFJochException(JFJochExceptionCategory::HDF5,
                                      "Legacy image source file missing");

            AppendOrExtendSourceMapping(ret, layout_.legacy_files.at(file_id), "/entry/data/data",
                                        source_image, local_image, 1);
        }

        return ret;
    }

    throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
                          "Unsupported HDF5 file layout for source mapping");
}