//  ************************************************************************************************
//
//  BornAgain: simulate and fit reflection and scattering
//
//! @file      Device/IO/ZipUtil.cpp
//! @brief     Implements functions in namespace ZipUtil.
//!
//! @homepage  http://www.bornagainproject.org
//! @license   GNU General Public License v3 or higher (see COPYING)
//! @copyright Forschungszentrum Jülich GmbH 2023
//! @authors   Scientific Computing Group at MLZ (see CITATION, AUTHORS)
//
//  ************************************************************************************************

#include "Device/IO/ZipUtil.h"
#include "Base/Util/PathUtil.h"
#include "Base/Util/StringUtil.h"
#include <bzlib.h> // For Bzip2 decompression
#include <fstream>
#include <zlib.h> // For Gzip decompression

namespace {

const std::string GzipExtension = ".gz";
const std::string BzipExtension = ".bz2";

bool isGZipped(const std::string& fname)
{
    return Base::Path::hasExtension(fname, GzipExtension);
}

bool isBZipped(const std::string& fname)
{
    return Base::Path::hasExtension(fname, BzipExtension);
}

bool isCompressed(const std::string& fname)
{
    return isGZipped(fname) || isBZipped(fname);
}

bool isNativeBinary(const std::string& fname)
{
    const std::string ext = ZipUtil::uncompressedExtension(fname);
    return ext == ".tif" || ext == ".tiff";
}

bool isBinary(const std::string& fname)
{
    return isCompressed(fname) || isNativeBinary(fname);
}

} // namespace


std::string ZipUtil::uncompressedExtension(std::string fname)
{
    if (isGZipped(fname))
        fname = fname.substr(0, fname.size() - GzipExtension.size());
    else if (isBZipped(fname))
        fname = fname.substr(0, fname.size() - BzipExtension.size());

    return Base::String::to_lower(Base::Path::extension(fname));
}

std::stringstream ZipUtil::file2stream(const std::string& fname)
{
    if (!Base::Path::IsFileExists(fname))
        throw std::runtime_error("File does not exist: " + fname);

    // Open the input file
    std::ifstream input_stream;
    std::ios_base::openmode openmode = std::ios::in;
    if (isBinary(fname))
        openmode |= std::ios_base::binary;
    input_stream.open(Base::Path::osPath(fname), openmode);

    if (!input_stream.is_open())
        throw std::runtime_error("Cannot open file for reading: " + fname);
    if (!input_stream.good())
        throw std::runtime_error("File is not good, probably it is a directory: " + fname);

    // Read the file content into a string
    std::ostringstream raw_stream;
    raw_stream << input_stream.rdbuf();
    std::string file_content = raw_stream.str();

    // Decompressed content will be written to this stringstream
    std::stringstream ss;

    // Decompress if necessary
    if (::isGZipped(fname)) {
        // Gzip decompression
        z_stream zs = {};
        zs.next_in = reinterpret_cast<Bytef*>(const_cast<char*>(file_content.data()));
        zs.avail_in = file_content.size();

        if (inflateInit2(&zs, 15 + 32) != Z_OK) // 15 + 32 enables Gzip decoding
            throw std::runtime_error("Failed to initialize Gzip decompression");

        char buffer[4096];
        do {
            zs.next_out = reinterpret_cast<Bytef*>(buffer);
            zs.avail_out = sizeof(buffer);

            int ret = inflate(&zs, Z_NO_FLUSH);
            if (ret != Z_OK && ret != Z_STREAM_END) {
                inflateEnd(&zs);
                throw std::runtime_error("Gzip decompression failed");
            }

            ss.write(buffer, sizeof(buffer) - zs.avail_out);
        } while (zs.avail_out == 0);

        inflateEnd(&zs);

    } else if (::isBZipped(fname)) {
        // Bzip2 decompression
        bz_stream bzs = {};
        bzs.next_in = const_cast<char*>(file_content.data());
        bzs.avail_in = file_content.size();

        if (BZ2_bzDecompressInit(&bzs, 0, 0) != BZ_OK)
            throw std::runtime_error("Failed to initialize Bzip2 decompression");

        char buffer[4096];
        do {
            bzs.next_out = buffer;
            bzs.avail_out = sizeof(buffer);

            int ret = BZ2_bzDecompress(&bzs);
            if (ret != BZ_OK && ret != BZ_STREAM_END) {
                BZ2_bzDecompressEnd(&bzs);
                throw std::runtime_error("Bzip2 decompression failed");
            }

            ss.write(buffer, sizeof(buffer) - bzs.avail_out);
        } while (bzs.avail_out == 0);

        BZ2_bzDecompressEnd(&bzs);

    } else {
        // If not compressed, simply copy the file content
        ss.str(file_content);
    }

    return ss;
}

void ZipUtil::stream2file(const std::string& fname, std::stringstream& s)
{
    std::ofstream fout;
    std::ios_base::openmode openmode = std::ios::out;
    if (isBinary(fname))
        openmode |= std::ios_base::binary;

    fout.open(Base::Path::osPath(fname), openmode);

    if (!fout.is_open())
        throw std::runtime_error("Cannot open file for writing: " + fname);
    if (!fout.good())
        throw std::runtime_error("File is not good, probably it is a directory: " + fname);

    const std::string& content = s.str(); // Get the content from the stringstream

    if (isGZipped(fname)) {
        // Gzip compression
        z_stream zs = {};
        if (deflateInit2(&zs, Z_DEFAULT_COMPRESSION, Z_DEFLATED, 15 + 16, 8, Z_DEFAULT_STRATEGY)
            != Z_OK)
            throw std::runtime_error("Failed to initialize Gzip compression");

        zs.next_in = reinterpret_cast<Bytef*>(const_cast<char*>(content.data()));
        zs.avail_in = content.size();

        char buffer[4096];
        do {
            zs.next_out = reinterpret_cast<Bytef*>(buffer);
            zs.avail_out = sizeof(buffer);

            int ret = deflate(&zs, Z_FINISH);
            if (ret != Z_OK && ret != Z_STREAM_END) {
                deflateEnd(&zs);
                throw std::runtime_error("Gzip compression failed");
            }

            fout.write(buffer, sizeof(buffer) - zs.avail_out);
        } while (zs.avail_out == 0);

        deflateEnd(&zs);

    } else if (isBZipped(fname)) {
        // Bzip2 compression
        bz_stream bzs = {};
        if (BZ2_bzCompressInit(&bzs, 9, 0, 0) != BZ_OK) // 9 = max compression level
            throw std::runtime_error("Failed to initialize Bzip2 compression");

        bzs.next_in = const_cast<char*>(content.data());
        bzs.avail_in = content.size();

        char buffer[4096];
        do {
            bzs.next_out = buffer;
            bzs.avail_out = sizeof(buffer);

            int ret = BZ2_bzCompress(&bzs, BZ_FINISH);
            if (ret != BZ_RUN_OK && ret != BZ_FINISH_OK && ret != BZ_STREAM_END) {
                BZ2_bzCompressEnd(&bzs);
                throw std::runtime_error("Bzip2 compression failed");
            }

            fout.write(buffer, sizeof(buffer) - bzs.avail_out);
        } while (bzs.avail_out == 0);

        BZ2_bzCompressEnd(&bzs);

    } else {
        // Write uncompressed content
        fout.write(content.data(), content.size());
    }

    fout.close();
}