/*
 * Copyright 2024 KylinSoft Co., Ltd.
 *
 * This program is free software: you can redistribute it and/or modify it under
 * the terms of the GNU General Public License as published by the Free Software
 * Foundation, either version 3 of the License, or (at your option) any later
 * version.
 *
 * This program is distributed in the hope that it will be useful, but WITHOUT
 * ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS
 * FOR A PARTICULAR PURPOSE. See the GNU General Public License for more
 * details.
 *
 * You should have received a copy of the GNU General Public License along with
 * this program. If not, see <https://www.gnu.org/licenses/>.
 */

#include "imageloader.h"

#include <cstring>
#include <string>

namespace {

bool isValidDimension(int value) { return value > 0 && value <= 0x7FFFFFFF; }

bool unsignedAddOverflows(unsigned int a, unsigned int b) {
    return a > static_cast<unsigned int>(-1) - b;
}

bool isPng(const unsigned char *data, unsigned int length) {
    static const unsigned char signature[] = {0x89, 0x50, 0x4E, 0x47, 0x0D,
                                              0x0A, 0x1A, 0x0A};
    return length >= sizeof(signature) &&
           std::memcmp(data, signature, sizeof(signature)) == 0;
}

bool isJpeg(const unsigned char *data, unsigned int length) {
    return length >= 2 && data[0] == 0xFF && data[1] == 0xD8;
}

bool isWebp(const unsigned char *data, unsigned int length) {
    return length >= 12 && std::memcmp(data, "RIFF", 4) == 0 &&
           std::memcmp(data + 8, "WEBP", 4) == 0;
}

std::string detectFormat(const unsigned char *data, unsigned int length) {
    if (isPng(data, length)) {
        return "png";
    }
    if (isJpeg(data, length)) {
        return "jpeg";
    }
    if (isWebp(data, length)) {
        return "webp";
    }
    return "unknown";
}

bool parsePngDimensions(const unsigned char *data, unsigned int length, int &width,
                        int &height) {
    if (length < 24 || std::memcmp(data + 12, "IHDR", 4) != 0) {
        return false;
    }

    const unsigned int ihdrLength =
        (static_cast<unsigned int>(data[8]) << 24) |
        (static_cast<unsigned int>(data[9]) << 16) |
        (static_cast<unsigned int>(data[10]) << 8) |
        static_cast<unsigned int>(data[11]);
    if (ihdrLength != 13) {
        return false;
    }

    width = (data[16] << 24) | (data[17] << 16) | (data[18] << 8) | data[19];
    height = (data[20] << 24) | (data[21] << 16) | (data[22] << 8) | data[23];
    return isValidDimension(width) && isValidDimension(height);
}

bool parseJpegDimensions(const unsigned char *data, unsigned int length, int &width,
                         int &height) {
    unsigned int index = 2;
    while (index + 8 < length) {
        if (data[index] != 0xFF) {
            ++index;
            continue;
        }

        if (index + 3 >= length) {
            return false;
        }

        const unsigned char marker = data[index + 1];
        if (marker == 0xD8 || marker == 0xD9) {
            index += 2;
            continue;
        }
        if (marker == 0x01 || (marker >= 0xD0 && marker <= 0xD7)) {
            index += 2;
            continue;
        }

        const unsigned int segmentLength =
            (static_cast<unsigned int>(data[index + 2]) << 8) | data[index + 3];
        if (segmentLength < 2) {
            return false;
        }

        const bool isStartOfFrame =
            (marker >= 0xC0 && marker <= 0xC3) ||
            (marker >= 0xC5 && marker <= 0xC7) ||
            (marker >= 0xC9 && marker <= 0xCB) ||
            (marker >= 0xCD && marker <= 0xCF);
        if (isStartOfFrame) {
            if (index + 8 >= length || segmentLength < 5) {
                return false;
            }
            height = (data[index + 5] << 8) | data[index + 6];
            width = (data[index + 7] << 8) | data[index + 8];
            return isValidDimension(width) && isValidDimension(height);
        }

        if (unsignedAddOverflows(index + 2U, segmentLength)) {
            return false;
        }
        const unsigned int nextIndex = index + 2 + segmentLength;
        if (nextIndex <= index + 2 || nextIndex > length) {
            return false;
        }
        index = nextIndex;
    }
    return false;
}

bool parseWebpDimensions(const unsigned char *data, unsigned int length, int &width,
                         int &height) {
    unsigned int index = 12;
    while (index + 8 <= length) {
        const unsigned char *chunk = data + index;
        const unsigned int chunkSize =
            chunk[4] | (chunk[5] << 8) | (chunk[6] << 16) | (chunk[7] << 24);
        if (unsignedAddOverflows(index + 8U, chunkSize) ||
            index + 8 + chunkSize > length) {
            break;
        }

        const unsigned char *payload = chunk + 8;
        if (std::memcmp(chunk, "VP8X", 4) == 0 && chunkSize >= 10) {
            width = (payload[4] | (payload[5] << 8) | (payload[6] << 16)) + 1;
            height = (payload[7] | (payload[8] << 8) | (payload[9] << 16)) + 1;
            return isValidDimension(width) && isValidDimension(height);
        }
        if (std::memcmp(chunk, "VP8 ", 4) == 0 && chunkSize >= 10) {
            width = payload[6] | (payload[7] << 8);
            height = payload[8] | (payload[9] << 8);
            width &= 0x3FFF;
            height &= 0x3FFF;
            return isValidDimension(width) && isValidDimension(height);
        }
        if (std::memcmp(chunk, "VP8L", 4) == 0 && chunkSize >= 5 &&
            payload[0] == 0x2F) {
            const unsigned int bits =
                payload[1] | (payload[2] << 8) | (payload[3] << 16) |
                (payload[4] << 24);
            width = (bits & 0x3FFF) + 1;
            height = ((bits >> 14) & 0x3FFF) + 1;
            return isValidDimension(width) && isValidDimension(height);
        }

        const unsigned int paddedSize = chunkSize + (chunkSize & 1U);
        if (unsignedAddOverflows(index + 8U, paddedSize)) {
            break;
        }
        const unsigned int nextIndex = index + 8 + paddedSize;
        if (nextIndex <= index || nextIndex > length) {
            break;
        }
        index = nextIndex;
    }
    return false;
}

bool parseDimensions(const unsigned char *data, unsigned int length,
                     const std::string &format, int &width, int &height) {
    if (format == "png") {
        return parsePngDimensions(data, length, width, height);
    }
    if (format == "jpeg") {
        return parseJpegDimensions(data, length, width, height);
    }
    if (format == "webp") {
        return parseWebpDimensions(data, length, width, height);
    }
    return false;
}

} // namespace

vision::image_loader::ImageBaseInfo
vision::image_loader::getImageInfoFromData(unsigned char *data,
                                           unsigned int length) {
    if (data == nullptr || length == 0) {
        return ImageBaseInfo{"unknown", 0, 0};
    }

    const std::string format = detectFormat(data, length);
    int width = 0;
    int height = 0;
    parseDimensions(data, length, format, width, height);

    return ImageBaseInfo{format, width, height};
}
