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

#include <tiffio.h>
#include <tiffio.hxx>
#include <sstream>

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


void WriteTIFF(TIFF *tiff, const CompressedImage& image) {

    if (tiff == nullptr)
        throw JFJochException(JFJochExceptionCategory::TIFFGeneratorError, "TIFFStreamOpen error");

    std::vector<uint8_t> buffer;
    image.GetUncompressed(buffer);

    TIFFSetField(tiff, TIFFTAG_IMAGEWIDTH, image.GetWidth());  // set the width of the image
    TIFFSetField(tiff, TIFFTAG_IMAGELENGTH, image.GetHeight());    // set the height of the image

    switch (image.GetMode()) {
        case CompressedImageMode::RGB:
            TIFFSetField(tiff, TIFFTAG_BITSPERSAMPLE, 8);
            TIFFSetField(tiff, TIFFTAG_SAMPLESPERPIXEL, 3);
            TIFFSetField(tiff, TIFFTAG_PHOTOMETRIC, PHOTOMETRIC_RGB);
            TIFFSetField(tiff, TIFFTAG_SAMPLEFORMAT, SAMPLEFORMAT_UINT);
            break;
        case CompressedImageMode::Int8:
            TIFFSetField(tiff, TIFFTAG_BITSPERSAMPLE, 8);
            TIFFSetField(tiff, TIFFTAG_SAMPLESPERPIXEL, 1);
            TIFFSetField(tiff, TIFFTAG_SAMPLEFORMAT, SAMPLEFORMAT_INT);
            break;
        case CompressedImageMode::Int16:
            TIFFSetField(tiff, TIFFTAG_BITSPERSAMPLE, 16);
            TIFFSetField(tiff, TIFFTAG_SAMPLESPERPIXEL, 1);
            TIFFSetField(tiff, TIFFTAG_SAMPLEFORMAT, SAMPLEFORMAT_INT);
            break;
        case CompressedImageMode::Int32:
            TIFFSetField(tiff, TIFFTAG_BITSPERSAMPLE, 32);
            TIFFSetField(tiff, TIFFTAG_SAMPLESPERPIXEL, 1);
            TIFFSetField(tiff, TIFFTAG_SAMPLEFORMAT, SAMPLEFORMAT_INT);
            break;
        case CompressedImageMode::Uint8:
            TIFFSetField(tiff, TIFFTAG_BITSPERSAMPLE, 8);
            TIFFSetField(tiff, TIFFTAG_SAMPLESPERPIXEL, 1);
            TIFFSetField(tiff, TIFFTAG_SAMPLEFORMAT, SAMPLEFORMAT_UINT);
            break;
        case CompressedImageMode::Uint16:
            TIFFSetField(tiff, TIFFTAG_BITSPERSAMPLE, 16);
            TIFFSetField(tiff, TIFFTAG_SAMPLESPERPIXEL, 1);
            TIFFSetField(tiff, TIFFTAG_SAMPLEFORMAT, SAMPLEFORMAT_UINT);
            break;
        case CompressedImageMode::Uint32:
            TIFFSetField(tiff, TIFFTAG_BITSPERSAMPLE, 32);
            TIFFSetField(tiff, TIFFTAG_SAMPLESPERPIXEL, 1);
            TIFFSetField(tiff, TIFFTAG_SAMPLEFORMAT, SAMPLEFORMAT_UINT);
            break;
        case CompressedImageMode::Float16:
            TIFFSetField(tiff, TIFFTAG_BITSPERSAMPLE, 16);
            TIFFSetField(tiff, TIFFTAG_SAMPLEFORMAT, SAMPLEFORMAT_IEEEFP);  // Float format
            TIFFSetField(tiff, TIFFTAG_SAMPLESPERPIXEL, 1);
            break;
        case CompressedImageMode::Float32:
            TIFFSetField(tiff, TIFFTAG_BITSPERSAMPLE, 32);
            TIFFSetField(tiff, TIFFTAG_SAMPLEFORMAT, SAMPLEFORMAT_IEEEFP);  // Float format
            TIFFSetField(tiff, TIFFTAG_SAMPLESPERPIXEL, 1);
            break;
        case CompressedImageMode::Float64:
            TIFFSetField(tiff, TIFFTAG_BITSPERSAMPLE, 64);
            TIFFSetField(tiff, TIFFTAG_SAMPLEFORMAT, SAMPLEFORMAT_IEEEFP);  // Float format
            TIFFSetField(tiff, TIFFTAG_SAMPLESPERPIXEL, 1);
            break;
    }
    TIFFSetField(tiff, TIFFTAG_PLANARCONFIG, PLANARCONFIG_CONTIG);
    TIFFSetField(tiff, TIFFTAG_ORIENTATION, ORIENTATION_TOPLEFT);
    TIFFSetField(tiff, TIFFTAG_COMPRESSION, COMPRESSION_NONE); // set compression to LZW
    TIFFSetField(tiff, TIFFTAG_ROWSPERSTRIP, image.GetHeight());

    if (TIFFWriteEncodedStrip(tiff, 0, buffer.data(), image.GetUncompressedSize()) < 0)
        throw JFJochException(JFJochExceptionCategory::TIFFGeneratorError, "TIFFWriteEncodedStrip error");
}

std::string WriteTIFFToString(const CompressedImage& image) {
    std::stringstream os;

    TIFF *tiff = TIFFStreamOpen("x", (std::ostream *) &os);
    WriteTIFF(tiff, image);
    TIFFClose(tiff);

    return os.str();
}

void WriteTIFFToFile(const std::string &filename, const CompressedImage& image) {
    TIFF *tiff = TIFFOpen(filename.c_str(), "w");

    WriteTIFF(tiff, image);

    TIFFClose(tiff);
}

CompressedImage ReadTIFF(const std::string &s, std::vector<uint8_t> &buffer) {
    if (s.empty())
        throw JFJochException(JFJochExceptionCategory::TIFFGeneratorError, "No TIFF file provided");
    uint32_t rows_per_string = 0;

    uint16_t elem_size = 0;
    uint16_t format = 0;
    uint32_t cols = 0;
    uint32_t lines = 0;

    std::istringstream input_TIFF_stream(s);
    TIFF* tiff = TIFFStreamOpen("MemTIFF", &input_TIFF_stream);
    if (tiff == nullptr)
        throw JFJochException(JFJochExceptionCategory::TIFFGeneratorError,"Not a proper TIFF file");

    TIFFGetField(tiff, TIFFTAG_IMAGEWIDTH, &cols);  // get the width of the image
    TIFFGetField(tiff, TIFFTAG_IMAGELENGTH, &lines);    // get the height of the image
    TIFFGetField(tiff, TIFFTAG_BITSPERSAMPLE, &elem_size);    // get the size of the channels
    TIFFGetField(tiff, TIFFTAG_ROWSPERSTRIP, &rows_per_string);
    TIFFGetField(tiff, TIFFTAG_SAMPLEFORMAT, &format);


    if (cols * lines * elem_size <= 0)
        throw JFJochException(JFJochExceptionCategory::TIFFGeneratorError,"Size wrong");

    if (cols * elem_size != TIFFScanlineSize(tiff))
        throw JFJochException(JFJochExceptionCategory::TIFFGeneratorError, "TIFFScanlineSize mismatch");

    buffer.resize(cols * lines * elem_size);

    for (int i = 0; i < lines; i++) {
        if (TIFFReadScanline(tiff, buffer.data() + i * cols * elem_size, i, 0) < 0)
            throw JFJochException(JFJochExceptionCategory::TIFFGeneratorError, "TIFFReadScanline error");
    }

    TIFFClose(tiff);

    CompressedImageMode mode = CalcImageMode(elem_size / 8, (format == SAMPLEFORMAT_IEEEFP), (format == SAMPLEFORMAT_INT));
    return CompressedImage(buffer, cols, lines, mode);
}

std::vector<uint32_t> ReadTIFFFromString32(const std::string &s, uint32_t &cols, uint32_t &lines) {
    if (s.empty())
        throw JFJochException(JFJochExceptionCategory::TIFFGeneratorError, "No TIFF file provided");

    uint32_t rows_per_string = 0;

    std::vector<uint32_t> ret;

    uint16_t elem_size;

    std::istringstream input_TIFF_stream(s);
    TIFF* tiff = TIFFStreamOpen("MemTIFF", &input_TIFF_stream);
    if (tiff == nullptr)
        throw JFJochException(JFJochExceptionCategory::TIFFGeneratorError,"Not a proper TIFF file");

    TIFFGetField(tiff, TIFFTAG_IMAGEWIDTH, &cols);  // get the width of the image
    TIFFGetField(tiff, TIFFTAG_IMAGELENGTH, &lines);    // get the height of the image
    TIFFGetField(tiff, TIFFTAG_BITSPERSAMPLE, &elem_size);    // get the size of the channels
    TIFFGetField(tiff, TIFFTAG_ROWSPERSTRIP, &rows_per_string);

    if (elem_size != 32)
        throw JFJochException(JFJochExceptionCategory::TIFFGeneratorError, "Only 32-bit format supported");

    ret.resize(cols * lines);

    if (cols * sizeof(uint32_t) != TIFFScanlineSize(tiff))
        throw JFJochException(JFJochExceptionCategory::TIFFGeneratorError, "TIFFScanlineSize mismatch");

    for (int i = 0; i < lines; i++) {
        if (TIFFReadScanline(tiff, ret.data() + i * cols, i, 0) < 0)
            throw JFJochException(JFJochExceptionCategory::TIFFGeneratorError, "TIFFReadScanline error");
    }

    TIFFClose(tiff);
    return ret;
}

std::vector<uint16_t> ReadTIFFFromString16(const std::string &s, uint32_t &cols, uint32_t &lines) {
    if (s.empty())
        throw JFJochException(JFJochExceptionCategory::TIFFGeneratorError, "No TIFF file provided");

    uint32_t rows_per_string = 0;

    std::vector<uint16_t> ret;

    uint16_t elem_size;

    std::istringstream input_TIFF_stream(s);
    TIFF* tiff = TIFFStreamOpen("MemTIFF", &input_TIFF_stream);
    if (tiff == nullptr)
        throw JFJochException(JFJochExceptionCategory::TIFFGeneratorError,
            "TIFF format error");

    TIFFGetField(tiff, TIFFTAG_IMAGEWIDTH, &cols);  // get the width of the image
    TIFFGetField(tiff, TIFFTAG_IMAGELENGTH, &lines);    // get the height of the image
    TIFFGetField(tiff, TIFFTAG_BITSPERSAMPLE, &elem_size);    // get the size of the channels
    TIFFGetField(tiff, TIFFTAG_ROWSPERSTRIP, &rows_per_string);

    if (elem_size != 16)
        throw JFJochException(JFJochExceptionCategory::TIFFGeneratorError, "Only 16-bit format supported");

    ret.resize(cols * lines);

    if (cols * sizeof(uint16_t) != TIFFScanlineSize(tiff))
        throw JFJochException(JFJochExceptionCategory::TIFFGeneratorError, "TIFFScanlineSize mismatch");

    for (int i = 0; i < lines; i++) {
        if (TIFFReadScanline(tiff, ret.data() + i * cols, i, 0) < 0)
            throw JFJochException(JFJochExceptionCategory::TIFFGeneratorError, "TIFFReadScanline error");
    }

    TIFFClose(tiff);
    return ret;
}

void SupressTIFFErrors() {
    TIFFSetErrorHandler(nullptr);
}