// SPDX-License-Identifier: LGPL-3.0-or-other
// Copyright (C) 2021 Contributors to the SLS Detector Package
#include "HDF5DataFile.h"
#include "receiver_defs.h"

#include <iomanip>

namespace sls {

HDF5DataFile::HDF5DataFile(int index, std::mutex *hdf5Lib)
    : File(HDF5), index_(index), hdf5Lib_(hdf5Lib) {

    parameterNames_ = std::vector<std::string>{
        "frame number",
        "exp length or sub exposure time",
        "packets caught",
        "bunch id",
        "timestamp",
        "mod id",
        "row",
        "column",
        "reserved",
        "debug",
        "round robin number",
        "detector type",
        "detector header version",
        "packets caught bit mask",
    };
    H5::StrType strdatatype(H5::PredType::C_S1, sizeof(bitset_storage));
    parameterDataTypes_ = std::vector<H5::DataType>{
        H5::PredType::STD_U64LE, H5::PredType::STD_U32LE, H5::PredType::STD_U32LE,
        H5::PredType::STD_U64LE, H5::PredType::STD_U64LE, H5::PredType::STD_U16LE,
        H5::PredType::STD_U16LE, H5::PredType::STD_U16LE, H5::PredType::STD_U16LE,
        H5::PredType::STD_U32LE, H5::PredType::STD_U16LE, H5::PredType::STD_U8LE,
        H5::PredType::STD_U8LE,  strdatatype};
}

HDF5DataFile::~HDF5DataFile() { CloseFile(); }

std::string HDF5DataFile::GetFileName() const {
    return fileName_;
}

uint32_t HDF5DataFile::GetFilesInAcquisition() const {
    return numFilesInAcquisition_;
}

H5::DataType HDF5DataFile::GetPDataType() const { return dataType_; }

std::vector<std::string> HDF5DataFile::GetParameterNames() const {
    return parameterNames_;
}
std::vector<H5::DataType> HDF5DataFile::GetParameterDataTypes() const {
    return parameterDataTypes_;
}

void HDF5DataFile::CloseFile() {
    std::lock_guard<std::mutex> lock(*hdf5Lib_);
    try {
        H5::Exception::dontPrint(); // to handle errors
        if (fd_) {
            fd_->close();
            delete fd_;
            fd_ = nullptr;
        }
    } catch (const H5::Exception &error) {
        LOG(logERROR) << "Could not close data HDF5 handles of index "
                      << index_;
        error.printErrorStack();
    }
    if (dataSpace_) {
        delete dataSpace_;
        dataSpace_ = nullptr;
    }
    if (dataSet_) {
        delete dataSet_;
        dataSet_ = nullptr;
    }
    if (dataSpacePara_) {
        delete dataSpacePara_;
        dataSpacePara_ = nullptr;
    }
    for (auto it : dataSetPara_)
        delete it;
    dataSetPara_.clear();
}

void HDF5DataFile::CreateFirstHDF5DataFile(
    const std::string filePath, const std::string fileNamePrefix,
    const uint64_t fileIndex, const bool overWriteEnable, const bool silentMode,
    const int modulePos, const int numUnitsPerReadout,
    const uint32_t udpPortNumber, const uint32_t maxFramesPerFile,
    const uint64_t numImages, const uint32_t nPixelsX, const uint32_t nPixelsY,
    const uint32_t dynamicRange) {

    subFileIndex_ = 0;
    numFramesInFile_ = 0;
    extNumImages_ = numImages;
    numFilesInAcquisition_ = 0;

    maxFramesPerFile_ = maxFramesPerFile;
    numImages_ = numImages;
    nPixelsX_ = nPixelsX;
    nPixelsY_ = nPixelsY;
    dynamicRange_ = dynamicRange;

    filePath_ = filePath;
    fileNamePrefix_ = fileNamePrefix;
    fileIndex_ = fileIndex;
    overWriteEnable_ = overWriteEnable;
    silentMode_ = silentMode;
    detIndex_ = modulePos;
    numUnitsPerReadout_ = numUnitsPerReadout;
    udpPortNumber_ = udpPortNumber;

    switch (dynamicRange_) {
    case 12:
    case 16:
        dataType_ = H5::PredType::STD_U16LE;
        break;
    case 32:
        dataType_ = H5::PredType::STD_U32LE;
        break;
    default:
        dataType_ = H5::PredType::STD_U8LE;
        break;
    }

    CreateFile();
}

void HDF5DataFile::CreateFile() {
    numFramesInFile_ = 0;
    numFilesInAcquisition_++;

    std::ostringstream os;
    os << filePath_ << "/" << fileNamePrefix_ << "_d"
       << (detIndex_ * numUnitsPerReadout_ + index_) << "_f" << subFileIndex_
       << '_' << fileIndex_ << ".h5";
    fileName_ = os.str();

    std::lock_guard<std::mutex> lock(*hdf5Lib_);

    uint64_t framestosave =
        ((maxFramesPerFile_ == 0) ? numImages_ : // infinite images
             (((extNumImages_ - subFileIndex_) > maxFramesPerFile_)
                  ? // save up to maximum at a time
                  maxFramesPerFile_
                  : (extNumImages_ - subFileIndex_)));

    uint64_t nDimx = framestosave;
    uint32_t nDimy = nPixelsY_;
    uint32_t nDimz = ((dynamicRange_ == 4) ? (nPixelsX_ / 2) : nPixelsX_);

    try {
        H5::Exception::dontPrint(); // to handle errors

        // file
        H5::FileAccPropList fapl;
        fapl.setFcloseDegree(H5F_CLOSE_STRONG);
        fd_ = nullptr;
        if (!overWriteEnable_)
            fd_ = new H5::H5File(fileName_.c_str(), H5F_ACC_EXCL,
                             H5::FileCreatPropList::DEFAULT, fapl);
        else
            fd_ = new H5::H5File(fileName_.c_str(), H5F_ACC_TRUNC,
                             H5::FileCreatPropList::DEFAULT, fapl);

        // attributes - version
        double dValue = HDF5_WRITER_VERSION;
        H5::DataSpace dataspace_attr = H5::DataSpace(H5S_SCALAR);
        H5::Attribute attribute = fd_->createAttribute(
            "version", H5::PredType::NATIVE_DOUBLE, dataspace_attr);
        attribute.write(H5::PredType::NATIVE_DOUBLE, &dValue);

        // dataspace
        hsize_t srcdims[3] = {nDimx, nDimy, nDimz};
        hsize_t srcdimsmax[3] = {H5S_UNLIMITED, nDimy, nDimz};
        dataSpace_ = nullptr;
        dataSpace_ = new H5::DataSpace(3, srcdims, srcdimsmax);

        // dataset
        // fill value
        H5::DSetCreatPropList plist;
        int fill_value = -1;
        plist.setFillValue(dataType_, &fill_value);
        // always create chunked dataset as unlimited is only
        // supported with chunked layout
        hsize_t chunk_dims[3] = {MAX_CHUNKED_IMAGES, nDimy, nDimz};
        plist.setChunk(3, chunk_dims);
        dataSet_ = nullptr;
        dataSet_ = new H5::DataSet(fd_->createDataSet(
            DATASET_NAME, dataType_, *dataSpace_, plist));

        // create parameter datasets
        hsize_t dims[1] = {nDimx};
        hsize_t dimsmax[1] = {H5S_UNLIMITED};
        dataSpacePara_ = nullptr;
        dataSpacePara_ = new H5::DataSpace(1, dims, dimsmax);

        // always create chunked dataset as unlimited is only
        // supported with chunked layout
        H5::DSetCreatPropList paralist;
        hsize_t chunkpara_dims[3] = {MAX_CHUNKED_IMAGES};
        paralist.setChunk(1, chunkpara_dims);

        for (unsigned int i = 0; i < parameterNames_.size(); ++i) {
            H5::DataSet *ds = new H5::DataSet(fd_->createDataSet(
                parameterNames_[i].c_str(), parameterDataTypes_[i],
                *dataSpacePara_, paralist));
            dataSetPara_.push_back(ds);
        }
    } catch (const H5::Exception &error) {
        error.printErrorStack();
        CloseFile();
        throw RuntimeError("Could not create HDF5 handles in object " +
                                index_);
    }
    if (!silentMode_) {
        LOG(logINFO) << "[" << udpPortNumber_
                     << "]: HDF5 File created: " << fileName_;
    }
}

void HDF5DataFile::WriteToFile(char *buffer, const int buffersize,
                               const uint64_t currentFrameNumber,
                               const uint32_t numPacketsCaught) {

    // check if maxframesperfile = 0 for infinite
    if (maxFramesPerFile_ && (numFramesInFile_ >= maxFramesPerFile_)) {
        CloseFile();
        ++subFileIndex_;
        CreateFile();
    }
    ++numFramesInFile_;

    // extend dataset (when receiver start followed by many status starts
    // (jungfrau)))
    if (currentFrameNumber >= extNumImages_) {
        ExtendDataset();
    }

    WriteDataFile(currentFrameNumber, buffer + sizeof(sls_receiver_header));
    WriteParameterDatasets(currentFrameNumber, (sls_receiver_header *)(buffer));
}

void HDF5DataFile::Convert12to16Bit(uint16_t *dst, uint8_t *src) {
    for (int i = 0; i < EIGER_NUM_PIXELS; ++i) {
        *dst = (uint16_t)(*src++ & 0xFF);
        *dst++ |= (uint16_t)((*src & 0xF) << 8u);
        ++i;
        *dst = (uint16_t)((*src++ & 0xF0) >> 4u);
        *dst++ |= (uint16_t)((*src++ & 0xFF) << 4u);
    }
}

void HDF5DataFile::WriteDataFile(const uint64_t currentFrameNumber,
                                 char *buffer) {
    // expand 12 bit to 16 bits
    char *revBuffer = buffer;
    if (dynamicRange_ == 12) {
        revBuffer = (char *)malloc(EIGER_16_BIT_IMAGE_SIZE);
        if (revBuffer == nullptr) {
            throw RuntimeError("Could not allocate memory for 12 bit to "
                                    "16 bit conversion in object " +
                                    std::to_string(index_));
        }
        Convert12to16Bit((uint16_t *)revBuffer, (uint8_t *)buffer);
    }

    std::lock_guard<std::mutex> lock(*hdf5Lib_);

    uint64_t nDimx =
        ((maxFramesPerFile_ == 0) ? currentFrameNumber
                                  : currentFrameNumber % maxFramesPerFile_);
    uint32_t nDimy = nPixelsY_;
    uint32_t nDimz = ((dynamicRange_ == 4) ? (nPixelsX_ / 2) : nPixelsX_);

    hsize_t count[3] = {1, nDimy, nDimz};
    hsize_t start[3] = {nDimx, 0, 0};
    hsize_t dims2[2] = {nDimy, nDimz};
    try {
        H5::Exception::dontPrint(); // to handle errors

        dataSpace_->selectHyperslab(H5S_SELECT_SET, count, start);
        H5::DataSpace memspace(2, dims2);
        dataSet_->write(revBuffer, dataType_, memspace, *dataSpace_);
        memspace.close();
        if (dynamicRange_ == 12) {
            free(revBuffer);
        }
    } catch (const H5::Exception &error) {
        if (dynamicRange_ == 12) {
            free(revBuffer);
        }
        LOG(logERROR) << "Could not write to file in object " << index_;
        error.printErrorStack();
        throw RuntimeError("Could not write to file in object " +
                                std::to_string(index_));
    }
}

void HDF5DataFile::WriteParameterDatasets(const uint64_t currentFrameNumber,
                                          sls_receiver_header *rheader) {
    std::lock_guard<std::mutex> lock(*hdf5Lib_);

    uint64_t fnum =
        ((maxFramesPerFile_ == 0) ? currentFrameNumber
                                  : currentFrameNumber % maxFramesPerFile_);

    sls_detector_header header = rheader->detHeader;
    hsize_t count[1] = {1};
    hsize_t start[1] = {fnum};
    int i = 0;
    try {
        H5::Exception::dontPrint(); // to handle errors
        dataSpacePara_->selectHyperslab(H5S_SELECT_SET, count, start);
        H5::DataSpace memspace(H5S_SCALAR);
        dataSetPara_[0]->write(&header.frameNumber, parameterDataTypes_[0],
                               memspace, *dataSpacePara_);
        i = 1;
        dataSetPara_[1]->write(&header.expLength, parameterDataTypes_[1],
                               memspace, *dataSpacePara_);
        i = 2;
        dataSetPara_[2]->write(&header.packetNumber, parameterDataTypes_[2],
                               memspace, *dataSpacePara_);
        i = 3;
        dataSetPara_[3]->write(&header.bunchId, parameterDataTypes_[3],
                               memspace, *dataSpacePara_);
        i = 4;
        dataSetPara_[4]->write(&header.timestamp, parameterDataTypes_[4],
                               memspace, *dataSpacePara_);
        i = 5;
        dataSetPara_[5]->write(&header.modId, parameterDataTypes_[5], memspace,
                               *dataSpacePara_);
        i = 6;
        dataSetPara_[6]->write(&header.row, parameterDataTypes_[6], memspace,
                               *dataSpacePara_);
        i = 7;
        dataSetPara_[7]->write(&header.column, parameterDataTypes_[7], memspace,
                               *dataSpacePara_);
        i = 8;
        dataSetPara_[8]->write(&header.reserved, parameterDataTypes_[8],
                               memspace, *dataSpacePara_);
        i = 9;
        dataSetPara_[9]->write(&header.debug, parameterDataTypes_[9], memspace,
                               *dataSpacePara_);
        i = 10;
        dataSetPara_[10]->write(&header.roundRNumber, parameterDataTypes_[10],
                                memspace, *dataSpacePara_);
        i = 11;
        dataSetPara_[11]->write(&header.detType, parameterDataTypes_[11],
                                memspace, *dataSpacePara_);
        i = 12;
        dataSetPara_[12]->write(&header.version, parameterDataTypes_[12],
                                memspace, *dataSpacePara_);
        i = 13;

        // contiguous bitset
        if (sizeof(sls_bitset) == sizeof(bitset_storage)) {
            dataSetPara_[13]->write((char *)&(rheader->packetsMask),
                                    parameterDataTypes_[13], memspace,
                                    *dataSpacePara_);
        }

        // not contiguous bitset
        else {
            // get contiguous representation of bit mask
            bitset_storage storage;
            memset(storage, 0, sizeof(bitset_storage));
            sls_bitset bits = rheader->packetsMask;
            for (int i = 0; i < MAX_NUM_PACKETS; ++i)
                storage[i >> 3] |= (bits[i] << (i & 7));
            // write bitmask
            dataSetPara_[13]->write((char *)storage, parameterDataTypes_[13],
                                    memspace, *dataSpacePara_);
        }
        i = 14;
    } catch (const H5::Exception &error) {
        error.printErrorStack();
        throw RuntimeError(
            "Could not write parameters (index:" + std::to_string(i) +
            ") to file in object " + std::to_string(index_));
    }
}

void HDF5DataFile::ExtendDataset() {
    std::lock_guard<std::mutex> lock(*hdf5Lib_);

    try {
        H5::Exception::dontPrint(); // to handle errors

        hsize_t dims[3];
        dataSpace_->getSimpleExtentDims(dims);
        dims[0] += numImages_;

        dataSet_->extend(dims);
        delete dataSpace_;
        dataSpace_ = nullptr;
        dataSpace_ = new H5::DataSpace(dataSet_->getSpace());

        hsize_t dims_para[1] = {dims[0]};
        for (unsigned int i = 0; i < dataSetPara_.size(); ++i)
            dataSetPara_[i]->extend(dims_para);
        delete dataSpacePara_;
        dataSpacePara_ = nullptr;
        dataSpacePara_ = new H5::DataSpace(dataSetPara_[0]->getSpace());

    } catch (const H5::Exception &error) {
        error.printErrorStack();
        throw RuntimeError("Could not extend dataset in object " +
                                std::to_string(index_));
    }
    if (!silentMode_) {
        LOG(logINFO) << index_ << " Extending HDF5 dataset by " << extNumImages_
                     << ", Total x Dimension: " << (extNumImages_ + numImages_);
    }
    extNumImages_ += numImages_;
}

} // namespace sls