// SPDX-License-Identifier: MIT
// DICOM image format.
// Copyright (C) 2024 Artem Senichev <artemsen@gmail.com>

#include "../imageformat.hpp"

#include <algorithm>
#include <cstring>
#include <limits>
#include <utility>

class ImageFormatDicom : public ImageFormat {
public:
    ImageFormatDicom() noexcept
        : ImageFormat(Priority::Low, "dicom")
    {
    }

    static constexpr const uint8_t SIG_DATA[] = { 'D', 'I', 'C', 'M' };
    static constexpr const size_t SIG_OFFSET = 128;

    [[nodiscard]] ImagePtr decode(const Data& data) const override
    {
        if (!check_signature(data, SIG_DATA, SIG_OFFSET)) {
            return nullptr;
        }

        DataStream stream(data, SIG_OFFSET + sizeof(SIG_DATA));

        // get image description
        DicomImage dicom {};
        if (!get_image(stream, dicom) || dicom.spp != 1 /* monochrome */ ||
            dicom.bpp != 16 /* 2 bytes per pixel */) {
            return nullptr;
        }

        // allocate image and frame
        ImagePtr image = std::make_shared<Image>();
        image->frames.resize(1);
        Pixmap& pm = image->frames[0].pm;
        pm.create(Pixmap::RGB, dicom.width, dicom.height);

        // calculate min/max color value if not set yet
        if (dicom.px_max == 0 || dicom.px_max <= dicom.px_min) {
            dicom.px_min = std::numeric_limits<int16_t>::max();
            dicom.px_max = std::numeric_limits<int16_t>::min();

            const int16_t* pixel_data =
                reinterpret_cast<const int16_t*>(dicom.data);
            const size_t num_pixels =
                static_cast<size_t>(dicom.width) * dicom.height;

            for (size_t i = 0; i < num_pixels; ++i) {
                const int16_t color = pixel_data[i];
                dicom.px_min = std::min(dicom.px_min, color);
                dicom.px_max = std::max(dicom.px_max, color);
            }
        }

        // calculate coefficient for converting 16-bit color to 8-bit
        double pixel_coeff;
        if (dicom.px_max <= dicom.px_min) {
            pixel_coeff = 1.0;
        } else {
            pixel_coeff = 256.0 / (dicom.px_max - dicom.px_min);
        }

        // decode image
        const int16_t* ptr = reinterpret_cast<const int16_t*>(dicom.data);
        pm.foreach([&ptr, dicom, pixel_coeff](argb_t& pixel) {
            int16_t src_color = *ptr++;
            src_color -= dicom.px_min;
            src_color *= pixel_coeff;
            const uint8_t dst_color = static_cast<uint8_t>(std::clamp(
                src_color, static_cast<int16_t>(0), static_cast<int16_t>(255)));
            pixel = { argb_t::max, dst_color, dst_color, dst_color };
        });

        image->format = "DICOM";

        return image;
    }

private:
    /** Data input stream. */
    struct DataStream {
        DataStream(const Data& raw_data, size_t pos)
            : data(raw_data)
            , position(pos)
        {
        }

        /**
         * Read data from the stream.
         * @param bytes number of bytes to consume
         * @return pointer to the data or nullptr on End of stream
         */
        const uint8_t* consume(const size_t bytes)
        {
            const uint8_t* ptr = nullptr;
            const size_t end = position + bytes;

            if (end <= data.size) {
                ptr = data.data + position;
                position = end;
            }

            return ptr;
        }

        /**
         * Read data from the stream.
         * @param value buffer for data to read
         * @return false on End of stream
         */
        template <typename T> bool read(T& value)
        {
            const uint8_t* ptr = consume(sizeof(T));
            if (ptr) {
                value = *reinterpret_cast<const T*>(ptr);
            }
            return ptr;
        }

    private:
        const Data& data;
        uint64_t position;
    };

    // DICOM tags
    static constexpr const uint32_t TAG_SAMPLES_PER_PIXEL = 0x00280002;
    static constexpr const uint32_t TAG_ROWS = 0x00280010;
    static constexpr const uint32_t TAG_COLUMNS = 0x00280011;
    static constexpr const uint32_t TAG_BIT_ALLOCATED = 0x00280100;
    static constexpr const uint32_t TAG_SMALL_PIXEL_VAL = 0x00280106;
    static constexpr const uint32_t TAG_BIG_PIXEL_VAL = 0x00280107;
    static constexpr const uint32_t TAG_PIXEL_DATA = 0x7fe00010;

    // DICOM element value types
    enum ValueRepresentation : uint16_t {
        VR_AE = 'A' | ('E' << 8),
        VR_AS = 'A' | ('S' << 8),
        VR_AT = 'A' | ('T' << 8),
        VR_CS = 'C' | ('S' << 8),
        VR_DA = 'D' | ('A' << 8),
        VR_DS = 'D' | ('S' << 8),
        VR_DT = 'D' | ('T' << 8),
        VR_FD = 'F' | ('D' << 8),
        VR_FL = 'F' | ('L' << 8),
        VR_IS = 'I' | ('S' << 8),
        VR_LO = 'L' | ('O' << 8),
        VR_LT = 'L' | ('T' << 8),
        VR_PN = 'P' | ('N' << 8),
        VR_SH = 'S' | ('H' << 8),
        VR_SL = 'S' | ('L' << 8),
        VR_SS = 'S' | ('S' << 8),
        VR_ST = 'S' | ('T' << 8),
        VR_TM = 'T' | ('M' << 8),
        VR_UI = 'U' | ('I' << 8),
        VR_UL = 'U' | ('L' << 8),
        VR_US = 'U' | ('S' << 8),
        VR_UT = 'U' | ('T' << 8),
        VR_OB = 'O' | ('B' << 8),
        VR_OW = 'O' | ('W' << 8),
        VR_SQ = 'S' | ('Q' << 8),
        VR_UN = 'U' | ('N' << 8),
        VR_QQ = 'Q' | ('Q' << 8),
        VR_RT = 'R' | ('T' << 8),
    };

    // DICOM image description
    struct DicomImage {
        uint16_t spp;        ///< Samples per Pixel
        uint16_t bpp;        ///< Number of bits allocated for each pixel sample
        uint16_t width;      ///< Image width
        uint16_t height;     ///< Image height
        int16_t px_min;      ///< Min pixel value encountered in the image
        int16_t px_max;      ///< Max pixel value encountered in the image
        const uint8_t* data; ///< Image data
        size_t data_sz;      ///< Size of data in bytes
    };

    // DICOM element description
    struct Element {
        uint32_t tag;
        uint16_t vr;
        const uint8_t* data;
        uint32_t size;
    };

    /**
     * Read next data element from the stream.
     * @param stream binary stream
     * @param element output element description
     * @return false if no more elements int the stream
     */
    static bool next_element(DataStream& stream, Element& element)
    {
        // read tag
        if (!stream.read(element.tag)) {
            return false;
        }
        element.tag = (element.tag << 16) | (element.tag >> 16);

        // read value representation (type)
        if (!stream.read(element.vr)) {
            return false;
        }

        // get payload size
        uint16_t size;
        if (!stream.read(size)) {
            return false;
        }
        element.size = size;
        if (element.size == 0 &&
            (element.vr == VR_OB || element.vr == VR_OW ||
             element.vr == VR_SQ || element.vr == VR_UN ||
             element.vr == VR_UT)) {
            if (!stream.read(element.size)) {
                return false;
            }
        }

        // get payload data
        if (element.size == 0) {
            element.data = nullptr;
        } else {
            element.data = stream.consume(element.size);
            if (!element.data) {
                return false;
            }
        }

        return true;
    }

    /**
     * Get image description from the stream.
     * @param stream binary stream
     * @param image output image description
     * @return true if image description is valid
     */
    static bool get_image(DataStream& stream, DicomImage& image)
    {
        Element el;

        while (next_element(stream, el)) {
            if (!el.data) {
                continue;
            }
            if (el.tag == TAG_SAMPLES_PER_PIXEL && el.vr == VR_US) {
                image.spp = *reinterpret_cast<const uint16_t*>(el.data);
            } else if (el.tag == TAG_ROWS && el.vr == VR_US) {
                image.height = *reinterpret_cast<const uint16_t*>(el.data);
            } else if (el.tag == TAG_COLUMNS && el.vr == VR_US) {
                image.width = *reinterpret_cast<const uint16_t*>(el.data);
            } else if (el.tag == TAG_BIT_ALLOCATED && el.vr == VR_US) {
                image.bpp = *reinterpret_cast<const uint16_t*>(el.data);
            } else if (el.tag == TAG_SMALL_PIXEL_VAL && el.vr == VR_SS) {
                image.px_min = *reinterpret_cast<const int16_t*>(el.data);
            } else if (el.tag == TAG_BIG_PIXEL_VAL && el.vr == VR_SS) {
                image.px_max = *reinterpret_cast<const int16_t*>(el.data);
            } else if (el.tag == TAG_PIXEL_DATA && el.vr == VR_OW) {
                image.data = el.data;
                image.data_sz = el.size;
            }
        }

        return image.data && image.height && image.width &&
            std::cmp_equal(image.data_sz,
                           image.width * image.height * (image.bpp / 8));
    }
};

// register format in factory
static ImageFormatDicom format_dicom;
