// 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"

#ifdef _WIN32
#include <windows.h>
#else
#include <fcntl.h>
#include <unistd.h>
#endif

// Positional reads: pread() on POSIX, ReadFile() with an OVERLAPPED offset on Windows. Both take the
// offset as an argument instead of moving a shared file position, which is what lets every worker
// thread read through one handle at the same time.
HDF5ImageSource::RawFile::RawFile(const std::string &path) {
#ifdef _WIN32
    HANDLE h = CreateFileA(path.c_str(), GENERIC_READ, FILE_SHARE_READ | FILE_SHARE_WRITE, nullptr,
                           OPEN_EXISTING, FILE_ATTRIBUTE_NORMAL, nullptr);
    handle_ = (h == INVALID_HANDLE_VALUE) ? -1 : reinterpret_cast<intptr_t>(h);
#else
    handle_ = ::open(path.c_str(), O_RDONLY);
#endif
}

HDF5ImageSource::RawFile::~RawFile() {
    if (handle_ == -1)
        return;
#ifdef _WIN32
    CloseHandle(reinterpret_cast<HANDLE>(handle_));
#else
    ::close(static_cast<int>(handle_));
#endif
}

void HDF5ImageSource::RawFile::ReadAt(void *dst, size_t size, uint64_t address) const {
    auto *out = static_cast<uint8_t *>(dst);
    size_t done = 0;
    while (done < size) {
#ifdef _WIN32
        OVERLAPPED ov{};
        ov.Offset = static_cast<DWORD>((address + done) & 0xFFFFFFFFULL);
        ov.OffsetHigh = static_cast<DWORD>((address + done) >> 32);
        DWORD got = 0;
        const bool ok = ReadFile(reinterpret_cast<HANDLE>(handle_), out + done,
                                 static_cast<DWORD>(size - done), &got, &ov);
        const long long n = ok ? static_cast<long long>(got) : -1;
#else
        const long long n = ::pread(static_cast<int>(handle_), out + done, size - done, address + done);
#endif
        if (n <= 0)
            throw JFJochException(JFJochExceptionCategory::HDF5, "Error reading image chunk from file");
        done += static_cast<size_t>(n);
    }
}

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

void HDF5ImageSource::Clear() {
    dataset_cache_.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,
                                                                     uint64_t stride) const {
    return locator_.GetSourceMapping(first_image, image_count, total_images, stride);
}

const HDF5ImageSource::OpenDataset &
HDF5ImageSource::GetDataset(const HDF5ImageLocator::Location &loc) const {
    if (auto it = dataset_cache_.find(loc.file.get()); it != dataset_cache_.end())
        return it->second;

    OpenDataset entry;
    entry.file = loc.file;
    entry.dataset = std::make_unique<HDF5DataSet>(*loc.file, "/entry/data/data");

    HDF5DataSpace dataspace(*entry.dataset);
    HDF5DataType datatype(*entry.dataset);
    HDF5Dcpl dcpl(*entry.dataset);

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

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

    auto dim = dataspace.GetDimensions();
    entry.height = dim[1];
    entry.width = dim[2];
    entry.mode = CalcImageMode(datatype.GetElemSize(), datatype.IsFloat(), datatype.IsSigned());

    auto chunk_size = dcpl.GetChunking();
    entry.direct_chunk = (chunk_size.size() == 3) && (chunk_size[0] == 1)
                         && (chunk_size[1] == dim[1]) && (chunk_size[2] == dim[2]);
    if (entry.direct_chunk)
        entry.algorithm = dcpl.GetCompression();

    if (entry.direct_chunk && !loc.path.empty()) {
        entry.raw = std::make_unique<RawFile>(loc.path);
        if (!entry.raw->IsOpen())
            entry.raw.reset();
        hid_t fcpl = H5Fget_create_plist(loc.file->GetID());
        if (fcpl >= 0) {
            hsize_t user_block = 0;
            if (H5Pget_userblock(fcpl, &user_block) >= 0)
                entry.user_block = user_block;
            H5Pclose(fcpl);
        }
    }

    return dataset_cache_.emplace(loc.file.get(), std::move(entry)).first->second;
}

std::optional<HDF5ImageSource::DirectChunk>
HDF5ImageSource::PrepareDirectRead(const HDF5ImageLocator::Location &loc) const {
    const auto &ds = GetDataset(loc);
    if (!ds.raw)
        return {};

    const hsize_t coord[3] = {static_cast<hsize_t>(loc.local_index), 0, 0};
    unsigned filter_mask = 0;
    haddr_t address = HADDR_UNDEF;
    hsize_t size = 0;
    if (H5Dget_chunk_info_by_coord(ds.dataset->GetID(), coord, &filter_mask, &address, &size) < 0)
        return {};
    // A chunk nobody ever wrote has no address and no bytes; only HDF5 knows it reads as the fill
    // value, so hand those back to it.
    if (address == HADDR_UNDEF || size == 0)
        return {};

    return DirectChunk{ds.raw.get(), ds.user_block + address, static_cast<uint32_t>(size),
                       ds.width, ds.height, ds.mode, ds.algorithm};
}

CompressedImage HDF5ImageSource::ReadDirect(std::vector<uint8_t> &buffer, const DirectChunk &chunk) {
    buffer.resize(chunk.size);
    chunk.file->ReadAt(buffer.data(), chunk.size, chunk.address);
    return {buffer, chunk.width, chunk.height, chunk.mode, chunk.algorithm};
}

CompressedImage HDF5ImageSource::ReadImageAt(std::vector<uint8_t> &buffer,
                                             const HDF5ImageLocator::Location &loc) const {
    const auto &ds = GetDataset(loc);
    const std::vector<hsize_t> start = {static_cast<hsize_t>(loc.local_index), 0, 0};

    if (ds.direct_chunk)
        ds.dataset->ReadDirectChunk(buffer, start);
    else
        ds.dataset->ReadVectorToU8(buffer, start, {1, ds.height, ds.width});

    return {buffer, ds.width, ds.height, ds.mode, ds.algorithm};
}